Add default_units and set_units on Snap

plonk/config.toml

@@ -34,7 +34,7 @@ accretion_radius = "length: 1"
 alpha_viscosity_numerical = ""
 ambipolar_diffusion_coefficient = "length: 2, time: -1"
 angular_momentum = "mass: 1, length: 2, time: -1"
-angular_velocity = "length: 1, time: -1"
+angular_velocity = "time: -1"
 azimuthal_angle = "angle: 1"
 density = "mass: 1, length: -3"
 differential_velocity = "length: 1, time: -1"
@@ -51,6 +51,7 @@ gravitational_potential = "mass: 1, length: 2, time: -2"
 grain_density = "mass: 1, length: -3"
 grain_size = "length: 1"
 hall_effect_coefficient = "length: 2, time: -1"
+id = ""
 inclination = "angle: 1"
 internal_energy = "length: 2, time: -2"
 keplerian_frequency = "time: -1"

plonk/snap/readers/phantom.py

@@ -10,7 +10,7 @@
 
 from ..._config import load_config
 from ..._logging import logger
-from ..._units import Quantity, generate_array_code_units
+from ..._units import Quantity, array_units, generate_array_code_units
 from ..._units import units as plonk_units
 from ..extra import extra_quantities
 from ..snap import Snap
@@ -79,7 +79,10 @@ def generate_snap_from_file(
     snap._properties, snap._code_units = header_to_properties(header)
     snap._array_code_units = generate_array_code_units(snap._code_units)
 
-    # OPTIONAL: Set snap._name_map
+    # OPTIONAL: Set snap._units.
+    snap._default_units = array_units(filename=config)
+
+    # OPTIONAL: Set snap._name_map.
     conf = load_config(filename=config)
     snap._name_map = {
         'particles': conf['phantom']['particles']['namemap'],

plonk/snap/snap.py

@@ -123,6 +123,7 @@ def __init__(self):
         self.data_source = None
         self.file_path = None
         self._code_units = {}
+        self._default_units = {}
         self._properties = {}
         self._array_code_units = {}
         self._array_registry: Dict[str, Callable] = {}
@@ -311,6 +312,7 @@ def bulk_load(self, arrays: List[str] = None) -> Snap:
             _arrays = arrays
         with self.context(cache=True):
             for array in _arrays:
+                print(array)
                 try:
                     self[array]
                 except ValueError as e:
@@ -344,6 +346,36 @@ def properties(self) -> Dict[str, Any]:
         """Snap properties."""
         return {key: self._properties[key] for key in sorted(self._properties.keys())}
 
+    @property
+    def default_units(self) -> Dict[str, Any]:
+        """Snap default units."""
+        return {
+            key: self._default_units[key] for key in sorted(self._default_units.keys())
+        }
+
+    def set_units(self, **kwargs) -> Snap:
+        """Set default unit for arrays.
+
+        Parameters
+        ----------
+        kwargs
+            Keyword arguments with keys as the array name, e.g.
+            'pressure', and with values as the unit as a string, e.g.
+            'pascal'.
+
+        Examples
+        --------
+        Set multiple default units.
+
+        >>> snap.set_units(pressure='pascal', density='g/cm^3')
+        """
+        for key, val in kwargs.items():
+            if key not in self.default_units:
+                logger.info(f'adding array {key} to default_units dict')
+            self._default_units[key] = val
+
+        return self
+
     @property
     def code_units(self) -> Dict[str, Any]:
         """Snap code units."""
@@ -1085,9 +1117,14 @@ def _get_array(self, name: str, sinks: bool = False) -> Quantity:
         else:
             array_dict = self._arrays
         if name in array_dict:
+            if name in self.default_units:
+                return array_dict[name].to(self.default_units[name])
             return array_dict[name]
         if name in self._array_registry or name in self._sink_registry:
-            array = self._get_array_from_registry(name, sinks)
+            if name in self.default_units:
+                array = self._get_array_from_registry(name, sinks).to(
+                    self.default_units[name]
+                )
             if self.cache_arrays:
                 if sinks:
                     self._sink_arrays[name] = array