From ac0373d05a748b0808e641487c3ed2c4d1fa5e00 Mon Sep 17 00:00:00 2001 From: Casper Welzel Andersen Date: Fri, 8 Jan 2021 17:37:53 +0100 Subject: [PATCH 1/3] Rename CLI `calc` to `recalc` Since the command will _always_ recalculate fields if already present, this is more appropriate. --- aiida_optimade/cli/__init__.py | 2 +- .../cli/{cmd_calc.py => cmd_recalc.py} | 8 +++++--- tests/cli/{test_calc.py => test_recalc.py} | 20 +++++++++---------- 3 files changed, 16 insertions(+), 14 deletions(-) rename aiida_optimade/cli/{cmd_calc.py => cmd_recalc.py} (95%) rename tests/cli/{test_calc.py => test_recalc.py} (92%) diff --git a/aiida_optimade/cli/__init__.py b/aiida_optimade/cli/__init__.py index 4c95e0ba..9d21cb88 100644 --- a/aiida_optimade/cli/__init__.py +++ b/aiida_optimade/cli/__init__.py @@ -5,4 +5,4 @@ click_completion.init() # Import to populate sub commands -from aiida_optimade.cli import cmd_calc, cmd_init, cmd_run # noqa: E402,F401 +from aiida_optimade.cli import cmd_recalc, cmd_init, cmd_run # noqa: E402,F401 diff --git a/aiida_optimade/cli/cmd_calc.py b/aiida_optimade/cli/cmd_recalc.py similarity index 95% rename from aiida_optimade/cli/cmd_calc.py rename to aiida_optimade/cli/cmd_recalc.py index 0879119f..fd668f1e 100644 --- a/aiida_optimade/cli/cmd_calc.py +++ b/aiida_optimade/cli/cmd_recalc.py @@ -35,8 +35,8 @@ help="Suppress informational output.", ) @click.pass_obj -def calc(obj: dict, fields: Tuple[str], force_yes: bool, silent: bool): - """Calculate OPTIMADE fields in the AiiDA database.""" +def recalc(obj: dict, fields: Tuple[str], force_yes: bool, silent: bool): + """Recalculate OPTIMADE fields in the AiiDA database.""" from aiida import load_profile from aiida.cmdline.utils import echo @@ -130,7 +130,9 @@ def calc(obj: dict, fields: Tuple[str], force_yes: bool, silent: bool): except Exception as exc: # pylint: disable=broad-except from traceback import print_exc - LOGGER.error("Full exception from 'aiida-optimade calc' CLI:\n%s", print_exc()) + LOGGER.error( + "Full exception from 'aiida-optimade recalc' CLI:\n%s", print_exc() + ) echo.echo_critical( f"An exception happened while trying to initialize {profile!r}:\n{exc!r}" ) diff --git a/tests/cli/test_calc.py b/tests/cli/test_recalc.py similarity index 92% rename from tests/cli/test_calc.py rename to tests/cli/test_recalc.py index c10c81ac..0e68068a 100644 --- a/tests/cli/test_calc.py +++ b/tests/cli/test_recalc.py @@ -1,6 +1,6 @@ # pylint: disable=unused-argument,too-many-locals def test_calc_all_new(run_cli_command, aiida_profile, top_dir): - """Test `aiida-optimade -p profile_name calc` works for non-existent fields. + """Test `aiida-optimade -p profile_name recalc` works for non-existent fields. By "non-existent" the meaning is calculating fields that don't already exist for any Nodes. @@ -8,7 +8,7 @@ def test_calc_all_new(run_cli_command, aiida_profile, top_dir): from aiida import orm from aiida.tools.importexport import import_data - from aiida_optimade.cli import cmd_calc + from aiida_optimade.cli import cmd_recalc from aiida_optimade.translators.entities import AiidaEntityTranslator # Clear database and get initialized_structure_nodes.aiida @@ -53,7 +53,7 @@ def test_calc_all_new(run_cli_command, aiida_profile, top_dir): ) options = ["--force-yes"] + fields - result = run_cli_command(cmd_calc.calc, options) + result = run_cli_command(cmd_recalc.recalc, options) assert ( f"Fields found for {n_structure_data} Nodes." not in result.stdout @@ -85,11 +85,11 @@ def test_calc_all_new(run_cli_command, aiida_profile, top_dir): def test_calc(run_cli_command, aiida_profile, top_dir): - """Test `aiida-optimade -p profile_name calc` works.""" + """Test `aiida-optimade -p profile_name recalc` works.""" from aiida import orm from aiida.tools.importexport import import_data - from aiida_optimade.cli import cmd_calc + from aiida_optimade.cli import cmd_recalc from aiida_optimade.translators.entities import AiidaEntityTranslator # Clear database and get initialized_structure_nodes.aiida @@ -113,7 +113,7 @@ def test_calc(run_cli_command, aiida_profile, top_dir): ) options = ["--force-yes"] + fields - result = run_cli_command(cmd_calc.calc, options) + result = run_cli_command(cmd_recalc.recalc, options) assert f"Fields found for {n_structure_data} Nodes." in result.stdout, result.stdout assert ( @@ -143,11 +143,11 @@ def test_calc(run_cli_command, aiida_profile, top_dir): def test_calc_partially_init(run_cli_command, aiida_profile, top_dir): - """Test `aiida-optimade -p profile_name calc` works for a partially initalized DB""" + """Test `aiida-optimade -p profile_name recalc` works for a partially initalized DB""" from aiida import orm from aiida.tools.importexport import import_data - from aiida_optimade.cli import cmd_calc + from aiida_optimade.cli import cmd_recalc from aiida_optimade.translators.entities import AiidaEntityTranslator # Clear database and get initialized_structure_nodes.aiida @@ -187,7 +187,7 @@ def test_calc_partially_init(run_cli_command, aiida_profile, top_dir): # "elements" should not be found in 3 Nodes options = ["--force-yes", "elements"] - result = run_cli_command(cmd_calc.calc, options) + result = run_cli_command(cmd_recalc.recalc, options) assert f"Field found for {n_total_nodes - 3} Nodes." in result.stdout, result.stdout assert ( @@ -216,7 +216,7 @@ def test_calc_partially_init(run_cli_command, aiida_profile, top_dir): # from one Node above. # This will also test if "elements_ratios" will be calculated from a # Node where both it and "elements" were missing prior to the previous - # invocation of `aiida-optimade calc`. + # invocation of `aiida-optimade recalc`. n_structure_data = ( orm.QueryBuilder() .append( From a4f69b78dca179d76dc9774268100c37f278b3ae Mon Sep 17 00:00:00 2001 From: Casper Welzel Andersen Date: Sat, 9 Jan 2021 01:01:09 +0100 Subject: [PATCH 2/3] Introduce initialization with minimization New `--minimized-fields` option for `aiida-optimade init`. This will try to avoid calculating big-valued OPTIONAL fields. To ensure this, all translator field methods have gotten a new property `store` on whether the calculated field should be stored in the `optimade` extra or not. This is to be able to calculate some fields that depend on others (like `structure_features`), but still _only_ saving the value of `structure_features` and not the others. Specifically for `structure_features` a post-treatment is needed to ensure the correct features are set/unset. --- aiida_optimade/cli/cmd_init.py | 30 ++++- aiida_optimade/entry_collections.py | 20 +++- aiida_optimade/mappers/structures.py | 18 +++ aiida_optimade/translators/entities.py | 12 +- aiida_optimade/translators/structures.py | 141 ++++++++++++++--------- 5 files changed, 149 insertions(+), 72 deletions(-) diff --git a/aiida_optimade/cli/cmd_init.py b/aiida_optimade/cli/cmd_init.py index 379af78e..d8dd1495 100644 --- a/aiida_optimade/cli/cmd_init.py +++ b/aiida_optimade/cli/cmd_init.py @@ -23,8 +23,19 @@ show_default=True, help="Suppress informational output.", ) +@click.option( + "-m", + "--minimized-fields", + is_flag=True, + default=False, + show_default=True, + help=( + "Do not calculate large-valued fields. This is especially good for structure " + "with thousands of atoms." + ), +) @click.pass_obj -def init(obj: dict, force: bool, silent: bool): +def init(obj: dict, force: bool, silent: bool, minimized_fields: bool): """Initialize an AiiDA database to be served with AiiDA-OPTIMADE.""" from aiida import load_profile from aiida.cmdline.utils import echo @@ -79,8 +90,21 @@ def init(obj: dict, force: bool, silent: bool): echo.echo_warning("This may take several minutes!") STRUCTURES._filter_fields = set() - STRUCTURES._alias_filter({"nelements": "2"}) - updated_pks = STRUCTURES._check_and_calculate_entities(cli=not silent) + if minimized_fields: + STRUCTURES._alias_filter( + dict.fromkeys( + [ + "structure_features", # required (will create species) + ], + None, + ) + ) + else: + STRUCTURES._alias_filter({"nsites": None}) + + updated_pks = STRUCTURES._check_and_calculate_entities( + cli=not silent, all_fields=not minimized_fields + ) except Exception as exc: # pylint: disable=broad-except from traceback import print_exc diff --git a/aiida_optimade/entry_collections.py b/aiida_optimade/entry_collections.py index 0b8fd856..ab7b3498 100644 --- a/aiida_optimade/entry_collections.py +++ b/aiida_optimade/entry_collections.py @@ -314,7 +314,9 @@ def _get_extras_filter_fields(self) -> set: if field.startswith(self.resource_mapper.PROJECT_PREFIX) } - def _check_and_calculate_entities(self, cli: bool = False) -> List[int]: + def _check_and_calculate_entities( + self, cli: bool = False, all_fields: bool = True + ) -> List[int]: """Check all entities have OPTIMADE extras, else calculate them For a bit of optimization, we only care about a field if it has specifically @@ -322,6 +324,8 @@ def _check_and_calculate_entities(self, cli: bool = False) -> List[int]: Parameters: cli: Whether or not this method is run through the CLI. + all_fields: Whether or not to calculate _all_ OPTIMADE fields or only those + defined through `filter_fields`. Returns: A list of the Node PKs representing the Nodes that were necessary to @@ -358,11 +362,17 @@ def _update_entities(entities: list, fields: list): necessary_entity_ids = [pk[0] for pk in necessary_entities_qb] # Create the missing OPTIMADE fields: - # All OPTIMADE fields fields = {"id", "type"} - fields |= self.get_attribute_fields() - # All provider-specific fields - fields |= {f"_{self.provider}_" + _ for _ in self.provider_fields} + if all_fields: + # All OPTIMADE fields + fields |= self.get_attribute_fields() + # All provider-specific fields + fields |= {f"_{self.provider}_" + _ for _ in self.provider_fields} + else: + # Only calculate for `filter_fields` + # "id" and "type" are ALWAYS needed though, hence `fields` is initiated + # with these values + fields |= self._get_extras_filter_fields() fields = list({self.resource_mapper.alias_for(f) for f in fields}) entities = self._find_all( diff --git a/aiida_optimade/mappers/structures.py b/aiida_optimade/mappers/structures.py index cfe7d4af..87f41e91 100644 --- a/aiida_optimade/mappers/structures.py +++ b/aiida_optimade/mappers/structures.py @@ -29,6 +29,7 @@ class StructureMapper(ResourceMapper): REQUIRED_ATTRIBUTES = set(StructureResourceAttributes.schema().get("required")) # This should be REQUIRED_FIELDS, but should be set as such in `optimade` + # pylint: disable=too-many-locals @classmethod def build_attributes( cls, retrieved_attributes: dict, entry_pk: int, node_type: str @@ -79,6 +80,23 @@ def build_attributes( ) else: res[attribute] = create_attribute() + # Special post-treatment for `structure_features` + all_fields = ( + translator._get_optimade_extras() + ) # pylint: disable=protected-access + all_fields.update(translator.new_attributes) + structure_features = all_fields.get("structure_features", []) + if all_fields.get("species", None) is None: + for feature in ["disorder", "implicit_atoms", "site_attachments"]: + try: + structure_features.remove(feature) + except ValueError: + # Not in list + pass + if structure_features != all_fields.get("structure_features", []): + # Some fields were removed + translator.new_attributes["structure_features"] = structure_features + # Store new attributes in `extras` translator.store_attributes() del translator diff --git a/aiida_optimade/translators/entities.py b/aiida_optimade/translators/entities.py index be19a038..35cc7958 100644 --- a/aiida_optimade/translators/entities.py +++ b/aiida_optimade/translators/entities.py @@ -54,19 +54,17 @@ def _node(self, value: Union[None, Node]): def _node_loaded(self): return bool(self.__node) - def _get_optimade_extras(self) -> Union[None, dict]: + def _get_optimade_extras(self) -> dict: if self._node_loaded: - return self._node.extras.get(self.EXTRAS_KEY, None) - return self._get_unique_node_property(f"extras.{self.EXTRAS_KEY}") + return self._node.extras.get(self.EXTRAS_KEY, {}) + res = self._get_unique_node_property(f"extras.{self.EXTRAS_KEY}") + return res or {} def store_attributes(self): """Store new attributes in Node extras and reset self._node""" if self.new_attributes: optimade = self._get_optimade_extras() - if optimade: - optimade.update(self.new_attributes) - else: - optimade = self.new_attributes + optimade.update(self.new_attributes) extras = ( self._get_unique_node_property("extras") if self._get_unique_node_property("extras") diff --git a/aiida_optimade/translators/structures.py b/aiida_optimade/translators/structures.py index f6948d3a..1e2a0309 100644 --- a/aiida_optimade/translators/structures.py +++ b/aiida_optimade/translators/structures.py @@ -137,7 +137,7 @@ def has_partial_occupancy(self) -> bool: return False # Start creating fields - def elements(self) -> List[str]: + def elements(self, store: bool = True) -> List[str]: """Names of elements found in the structure as a list of strings, in alphabetical order.""" attribute = "elements" @@ -151,23 +151,25 @@ def elements(self) -> List[str]: res.remove("X") # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def nelements(self) -> int: + def nelements(self, store: bool = True) -> int: """Number of different elements in the structure as an integer.""" attribute = "nelements" if attribute in self.new_attributes: return self.new_attributes[attribute] - res = len(self.elements()) + res = len(self.elements(store=store)) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def elements_ratios(self) -> List[float]: + def elements_ratios(self, store: bool = True) -> List[float]: """Relative proportions of different elements in the structure.""" attribute = "elements_ratios" @@ -177,13 +179,14 @@ def elements_ratios(self) -> List[float]: ratios = self.get_symbol_weights() total_weight = fsum(ratios.values()) - res = [ratios[symbol] / total_weight for symbol in self.elements()] + res = [ratios[symbol] / total_weight for symbol in self.elements(store=store)] # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = floats_to_hex(res) + if store: + self.new_attributes[attribute] = floats_to_hex(res) return res - def chemical_formula_descriptive(self) -> str: + def chemical_formula_descriptive(self, store: bool = True) -> str: """The chemical formula for a structure as a string in a form chosen by the API implementation.""" attribute = "chemical_formula_descriptive" @@ -193,10 +196,11 @@ def chemical_formula_descriptive(self) -> str: res = self.get_formula() # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def chemical_formula_reduced(self) -> str: + def chemical_formula_reduced(self, store: bool = True) -> str: """The reduced chemical formula for a structure As a string with element symbols and integer chemical proportion numbers. @@ -228,13 +232,16 @@ def chemical_formula_reduced(self) -> str: occupation[symbol] = "" else: occupation[symbol] = rounded_weight - res = "".join([f"{symbol}{occupation[symbol]}" for symbol in self.elements()]) + res = "".join( + [f"{symbol}{occupation[symbol]}" for symbol in self.elements(store=store)] + ) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def chemical_formula_hill(self) -> str: + def chemical_formula_hill(self, store: bool = True) -> str: """The chemical formula for a structure in Hill form With element symbols followed by integer chemical proportion numbers. @@ -256,10 +263,11 @@ def chemical_formula_hill(self) -> str: res = self.get_formula(mode="hill") # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def chemical_formula_anonymous(self) -> str: + def chemical_formula_anonymous(self, store: bool = True) -> str: """The anonymous formula is the chemical_formula_reduced But where the elements are instead first ordered by their chemical proportion number, @@ -277,7 +285,7 @@ def chemical_formula_anonymous(self) -> str: assert len(ANONYMOUS_ELEMENTS) >= len( weights - ), f"Not enough generated anonymous elements to create `chemical_formula_anonymous` for Node . Found elements: {len(self.elements())}. Generated anonymous elements: {len(ANONYMOUS_ELEMENTS)}." + ), f"Not enough generated anonymous elements to create `chemical_formula_anonymous` for Node . Found elements: {self.nelements(store=False)}. Generated anonymous elements: {len(ANONYMOUS_ELEMENTS)}." res = "" for index, occupation in enumerate(sorted(weights, reverse=True)): @@ -285,10 +293,11 @@ def chemical_formula_anonymous(self) -> str: res += f"{ANONYMOUS_ELEMENTS[index]}{rounded_weight}" # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def dimension_types(self) -> List[int]: + def dimension_types(self, store: bool = True) -> List[int]: """List of three integers. For each of the three directions indicated by the three lattice vectors @@ -304,10 +313,11 @@ def dimension_types(self) -> List[int]: res = self._pbc # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def nperiodic_dimensions(self) -> int: + def nperiodic_dimensions(self, store: bool = True) -> int: """Number of periodic dimensions.""" attribute = "nperiodic_dimensions" @@ -317,10 +327,11 @@ def nperiodic_dimensions(self) -> int: res = sum(self._pbc) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def lattice_vectors(self) -> List[List[float]]: + def lattice_vectors(self, store: bool = True) -> List[List[float]]: """The three lattice vectors in Cartesian coordinates, in ångström (Å).""" attribute = "lattice_vectors" @@ -330,10 +341,13 @@ def lattice_vectors(self) -> List[List[float]]: res = check_floating_round_errors(self._cell) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = floats_to_hex(res) + if store: + self.new_attributes[attribute] = floats_to_hex(res) return res - def cartesian_site_positions(self) -> List[List[Union[float, None]]]: + def cartesian_site_positions( + self, store: bool = True + ) -> List[List[Union[float, None]]]: """Cartesian positions of each site. A site is an atom, a site potentially occupied by an atom, @@ -348,23 +362,25 @@ def cartesian_site_positions(self) -> List[List[Union[float, None]]]: res = check_floating_round_errors(sites) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = floats_to_hex(res) + if store: + self.new_attributes[attribute] = floats_to_hex(res) return res - def nsites(self) -> int: + def nsites(self, store: bool = True) -> int: """An integer specifying the length of the cartesian_site_positions property.""" attribute = "nsites" if attribute in self.new_attributes: return self.new_attributes[attribute] - res = len(self.cartesian_site_positions()) + res = len(self.cartesian_site_positions(store=store)) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def species_at_sites(self) -> List[str]: + def species_at_sites(self, store: bool = True) -> List[str]: """Name of the species at each site (Where values for sites are specified with the same order of the property @@ -378,10 +394,11 @@ def species_at_sites(self) -> List[str]: res = [site["kind_name"] for site in self._sites] # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def species(self) -> List[dict]: + def species(self, store: bool = True) -> List[dict]: """A list describing the species of the sites of this structure. Species can be pure chemical elements, or virtual-crystal atoms @@ -429,10 +446,11 @@ def species(self) -> List[dict]: res.append(species) # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def assemblies(self) -> Union[List[dict], None]: + def assemblies(self, store: bool = True) -> Union[List[dict], None]: """A description of groups of sites that are statistically correlated. NOTE: Currently not supported. @@ -445,14 +463,12 @@ def assemblies(self) -> Union[List[dict], None]: res = None # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res - def structure_features(self) -> List[str]: - """A list of strings that flag which special features are used by the structure. - - SHOULD be absent if there are no partial occupancies - """ + def structure_features(self, store: bool = True) -> List[str]: + """A sorted list of strings that flag which special features are used by the structure.""" attribute = "structure_features" if attribute in self.new_attributes: @@ -460,15 +476,15 @@ def structure_features(self) -> List[str]: res = [] - # Figure out if there are partial occupancies - if not self.has_partial_occupancy(): - self.new_attributes[attribute] = res - return res + # * Assemblies * + # This flag MUST be present if the property assemblies is present. + if self.assemblies(store=False): + res.append("assemblies") # * Disorder * # This flag MUST be present if any one entry in the species list # has a chemical_symbols list that is longer than 1 element. - species = self.species() + species = self.species(store=False) key = "chemical_symbols" for item in species: if key not in item: @@ -479,20 +495,31 @@ def structure_features(self) -> List[str]: res.append("disorder") break - # * Unknown positions * - # This flag MUST be present if at least one component of the cartesian_site_positions - # list of lists has value null. - cartesian_site_positions = self.cartesian_site_positions() - for site in cartesian_site_positions: - if float("NaN") in site: - res.append("unknown_positions") + # * Implicit atoms * + # This flag MUST be present if the structure contains atoms that are not + # assigned to sites via the property species_at_sites (e.g., because their + # positions are unknown). When this flag is present, the properties related to + # the chemical formula will likely not match the type and count of atoms + # represented by the species_at_sites, species, and assemblies properties. + species_at_sites = self.species_at_sites(store=False) + key = "name" + for item in species: + if key not in item: + raise OptimadeIntegrityError( + f'The required key {key} was not found for {item} in the "species" attribute' + ) + if item[key] not in species_at_sites: + res.append("implicit_atoms") break - # * Assemblies * - # This flag MUST be present if the property assemblies is present. - if self.assemblies(): - res.append("assemblies") + # * Site attachements * + # This flag MUST be present if any one entry in the species list includes attached and nattached. + for item in species: + if item.get("attached", None) is not None: + res.append("site_attachments") + break # Finally, save OPTIMADE attribute for later storage in extras for AiiDA Node and return value - self.new_attributes[attribute] = res + if store: + self.new_attributes[attribute] = res return res From 88e7654281b73b85631af74600e4c505aeba65bf Mon Sep 17 00:00:00 2001 From: Casper Welzel Andersen Date: Sat, 9 Jan 2021 03:23:22 +0100 Subject: [PATCH 3/3] Ensure minimized fields works Build attributes only with supplied "desired" fields. This ensures the `-m/--minimized-fields` option for `aiida-optimade init` works as intended. --- aiida_optimade/cli/cmd_init.py | 17 ++++++----- aiida_optimade/entry_collections.py | 9 ++++-- aiida_optimade/mappers/entries.py | 12 +++++++- aiida_optimade/mappers/structures.py | 42 +++++++++++++++++++++------- aiida_optimade/models/structures.py | 11 +++++--- 5 files changed, 67 insertions(+), 24 deletions(-) diff --git a/aiida_optimade/cli/cmd_init.py b/aiida_optimade/cli/cmd_init.py index d8dd1495..70059061 100644 --- a/aiida_optimade/cli/cmd_init.py +++ b/aiida_optimade/cli/cmd_init.py @@ -91,14 +91,17 @@ def init(obj: dict, force: bool, silent: bool, minimized_fields: bool): STRUCTURES._filter_fields = set() if minimized_fields: - STRUCTURES._alias_filter( - dict.fromkeys( - [ - "structure_features", # required (will create species) - ], - None, - ) + minimized_keys = ( + STRUCTURES.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS.copy() + ) + minimized_keys |= STRUCTURES.get_attribute_fields() + minimized_keys |= { + f"_{STRUCTURES.provider}_" + _ for _ in STRUCTURES.provider_fields + } + minimized_keys.difference_update( + {"cartesian_site_positions", "nsites", "species_at_sites"} ) + STRUCTURES._alias_filter(dict.fromkeys(minimized_keys, None)) else: STRUCTURES._alias_filter({"nsites": None}) diff --git a/aiida_optimade/entry_collections.py b/aiida_optimade/entry_collections.py index ab7b3498..f7c283b9 100644 --- a/aiida_optimade/entry_collections.py +++ b/aiida_optimade/entry_collections.py @@ -153,7 +153,7 @@ def find( all_fields = criteria.pop("fields") if getattr(params, "response_fields", False): fields = set(params.response_fields.split(",")) - fields |= self.resource_mapper.get_required_fields() + fields |= self.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS else: fields = all_fields.copy() @@ -362,7 +362,7 @@ def _update_entities(entities: list, fields: list): necessary_entity_ids = [pk[0] for pk in necessary_entities_qb] # Create the missing OPTIMADE fields: - fields = {"id", "type"} + fields = self.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS.copy() if all_fields: # All OPTIMADE fields fields |= self.get_attribute_fields() @@ -373,6 +373,11 @@ def _update_entities(entities: list, fields: list): # "id" and "type" are ALWAYS needed though, hence `fields` is initiated # with these values fields |= self._get_extras_filter_fields() + fields |= { + f"_{self.provider}_" + _ + for _ in self._filter_fields + if _ in self.provider_fields + } fields = list({self.resource_mapper.alias_for(f) for f in fields}) entities = self._find_all( diff --git a/aiida_optimade/mappers/entries.py b/aiida_optimade/mappers/entries.py index e29b231c..63cb8ce0 100644 --- a/aiida_optimade/mappers/entries.py +++ b/aiida_optimade/mappers/entries.py @@ -16,7 +16,6 @@ class ResourceMapper(OptimadeResourceMapper): TRANSLATORS: Dict[str, AiidaEntityTranslator] ALL_ATTRIBUTES: set = set() - REQUIRED_ATTRIBUTES: set = set() @classmethod def all_aliases(cls) -> Tuple[Tuple[str, str]]: @@ -40,6 +39,8 @@ def map_back(cls, entity_properties: dict) -> dict: :return: A resource object in OPTIMADE format :rtype: dict """ + from optimade.server.config import CONFIG + new_object_attributes = {} new_object = {} @@ -64,8 +65,13 @@ def map_back(cls, entity_properties: dict) -> dict: if value is not None: new_object[field] = value + mapping = {aiida: optimade for optimade, aiida in cls.all_aliases()} + new_object["attributes"] = cls.build_attributes( retrieved_attributes=new_object_attributes, + desired_attributes={mapping.get(_, _) for _ in entity_properties} + - cls.TOP_LEVEL_NON_ATTRIBUTES_FIELDS + - set(CONFIG.aliases.get(cls.ENDPOINT, {}).keys()), entry_pk=new_object["id"], node_type=new_object["type"], ) @@ -77,6 +83,7 @@ def map_back(cls, entity_properties: dict) -> dict: def build_attributes( cls, retrieved_attributes: dict, + desired_attributes: list, entry_pk: int, node_type: str, ) -> dict: @@ -85,6 +92,9 @@ def build_attributes( :param retrieved_attributes: Dict of new attributes, will be updated accordingly :type retrieved_attributes: dict + :param desired_attributes: Set of attributes to be built. + :type desired_attributes: set + :param entry_pk: The AiiDA Node's PK :type entry_pk: int diff --git a/aiida_optimade/mappers/structures.py b/aiida_optimade/mappers/structures.py index 87f41e91..ca58ef76 100644 --- a/aiida_optimade/mappers/structures.py +++ b/aiida_optimade/mappers/structures.py @@ -26,21 +26,30 @@ class StructureMapper(ResourceMapper): "data.structure.StructureData.": StructureDataTranslator, } ALL_ATTRIBUTES = set(StructureResourceAttributes.schema().get("properties").keys()) - REQUIRED_ATTRIBUTES = set(StructureResourceAttributes.schema().get("required")) - # This should be REQUIRED_FIELDS, but should be set as such in `optimade` + REQUIRED_FIELDS = set(StructureResourceAttributes.schema().get("required")) # pylint: disable=too-many-locals @classmethod def build_attributes( - cls, retrieved_attributes: dict, entry_pk: int, node_type: str + cls, + retrieved_attributes: dict, + desired_attributes: set, + entry_pk: int, + node_type: str, ) -> dict: """Build attributes dictionary for OPTIMADE structure resource :param retrieved_attributes: Dict of new attributes, will be updated accordingly :type retrieved_attributes: dict + :param desired_attributes: List of attributes to be built. + :type desired_attributes: set + :param entry_pk: The AiiDA Node's PK :type entry_pk: int + + :param node_type: The AiiDA Node's type + :type node_type: str """ float_fields = { "elements_ratios", @@ -49,22 +58,30 @@ def build_attributes( } # Add existing attributes - missing_attributes = cls.ALL_ATTRIBUTES.copy() existing_attributes = set(retrieved_attributes.keys()) - missing_attributes.difference_update(existing_attributes) + desired_attributes.difference_update(existing_attributes) for field in float_fields: if field in existing_attributes and retrieved_attributes.get(field): retrieved_attributes[field] = hex_to_floats(retrieved_attributes[field]) res = retrieved_attributes.copy() + none_value_attributes = cls.REQUIRED_FIELDS - desired_attributes.union( + existing_attributes + ) + none_value_attributes = { + _ for _ in none_value_attributes if not _.startswith("_") + } + res.update({field: None for field in none_value_attributes}) + # Create and add new attributes - if missing_attributes: + if desired_attributes: translator = cls.TRANSLATORS[node_type](entry_pk) - for attribute in missing_attributes: + + for attribute in desired_attributes: try: create_attribute = getattr(translator, attribute) except AttributeError as exc: - if attribute in cls.REQUIRED_ATTRIBUTES: + if attribute in cls.get_required_fields(): translator = None raise NotImplementedError( f"Parsing required attribute {attribute!r} from " @@ -80,10 +97,11 @@ def build_attributes( ) else: res[attribute] = create_attribute() + # Special post-treatment for `structure_features` all_fields = ( - translator._get_optimade_extras() - ) # pylint: disable=protected-access + translator._get_optimade_extras() # pylint: disable=protected-access + ) all_fields.update(translator.new_attributes) structure_features = all_fields.get("structure_features", []) if all_fields.get("species", None) is None: @@ -97,6 +115,10 @@ def build_attributes( # Some fields were removed translator.new_attributes["structure_features"] = structure_features + translator.new_attributes.update( + {field: None for field in none_value_attributes} + ) + # Store new attributes in `extras` translator.store_attributes() del translator diff --git a/aiida_optimade/models/structures.py b/aiida_optimade/models/structures.py index 2d77a3a4..2342f765 100644 --- a/aiida_optimade/models/structures.py +++ b/aiida_optimade/models/structures.py @@ -1,12 +1,12 @@ # pylint: disable=missing-class-docstring,too-few-public-methods from datetime import datetime - -from pydantic import Field +from typing import Optional from optimade.models import ( StructureResource as OptimadeStructureResource, StructureResourceAttributes as OptimadeStructureResourceAttributes, ) +from optimade.models.utils import OptimadeField, SupportLevel def prefix_provider(string: str) -> str: @@ -21,8 +21,11 @@ def prefix_provider(string: str) -> str: class StructureResourceAttributes(OptimadeStructureResourceAttributes): """Extended StructureResourceAttributes for AiiDA-specific fields""" - ctime: datetime = Field( - ..., description="Creation time of the Node in the AiiDA database." + ctime: Optional[datetime] = OptimadeField( + ..., + description="Creation time of the Node in the AiiDA database.", + support=SupportLevel.SHOULD, + queryable=SupportLevel.MUST, ) class Config: