Build OPTIMADE fields according to supplied fields

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

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,24 @@ 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:
+            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})
+
+        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
 

aiida_optimade/cli/cmd_calc.py → 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}"
         )

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()
 
@@ -314,14 +314,18 @@ 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
         been queried for using "filter".
 
         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,22 @@ 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}
+            fields = self.resource_mapper.TOP_LEVEL_NON_ATTRIBUTES_FIELDS.copy()
+            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 |= {
+                    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(

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
 

aiida_optimade/mappers/structures.py

@@ -26,20 +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",
@@ -48,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 "
@@ -79,6 +97,28 @@ 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
+
+            translator.new_attributes.update(
+                {field: None for field in none_value_attributes}
+            )
+
             # Store new attributes in `extras`
             translator.store_attributes()
             del translator

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:

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")