Subdomain extraction

src/osyris/core/vector.py

@@ -209,7 +209,9 @@ def __invert__(self):
         return np.logical_not(self)
 
     def to(self, unit):
-        return self.__class__(**{c: xyz.to(unit) for c, xyz in self._xyz.items()})
+        return self.__class__(**{c: xyz.to(unit)
+                                 for c, xyz in self._xyz.items()},
+                              name=self.name)
 
     def _wrap_numpy(self, func, *args, **kwargs):
         if isinstance(args[0], (tuple, list)):

src/osyris/spatial/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2022 Osyris contributors (https://github.com/nvaytet/osyris)
+
+# flake8: noqa
+
+from .spatial import *

src/osyris/spatial/spatial.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2022 Osyris contributors (https://github.com/nvaytet/osyris)
+
+import numpy as np
+
+
+def extract_sphere(dataset, radius, origin):
+    """
+    Extract a spherical subdomain around an origin point.
+    """
+    subdomain = dataset.__class__(nout=dataset.meta["nout"], path=dataset.meta["path"])
+    subdomain.meta = dataset.meta
+    subdomain._parent = dataset
+
+    for name, group in dataset.items():
+        pos = group.get("position", group.parent["amr"]["position"])
+        r = (pos - origin).norm
+        c = (r < radius).values
+        if np.any(c):
+            subdomain[name] = dataset[name][c]
+
+    return subdomain
+
+
+def extract_cube(dataset, dx, dy, dz, origin):
+    """
+    Extract a cubic domain of size dx, dy & dz around an origin point
+    """
+    subdomain = dataset.__class__(nout=dataset.meta["nout"], path=dataset.meta["path"])
+    subdomain.meta = dataset.meta
+    subdomain._parent = dataset
+
+    for name, group in dataset.items():
+        pos = group.get("position", group.parent["amr"]["position"])
+        centered_pos = pos - origin
+        cx = (centered_pos.x <= dx * .5) & (centered_pos.x >= -dx * .5)
+        cy = (centered_pos.y <= dy * .5) & (centered_pos.y >= -dy * .5)
+        cz = (centered_pos.z <= dz * .5) & (centered_pos.z >= -dz * .5)
+        c = (cx & cy & cz).values
+        if np.any(c):
+            subdomain[name] = dataset[name][c]
+
+    return subdomain