Subdomain extraction

src/osyris/__init__.py

@@ -6,3 +6,4 @@
 from .config import config, units
 from .plot import histogram1d, histogram2d, plane, scatter, map, plot
 from .core import Array, Datagroup, Dataset, Plot
+from .spatial import spatial

src/osyris/core/array.py

@@ -47,7 +47,10 @@ def __getitem__(self, slice_):
                              unit=self._unit,
                              parent=self._parent,
                              name=self._name)
-        if isinstance(slice_, int) and self.ndim > 1:
+        if self.shape[0] == 1:
+            return out
+        if (isinstance(slice_, int)
+                or isinstance(slice_, np.integer)) and self.ndim > 1:
             return out.reshape(1, len(out))
         return out
 
@@ -253,6 +256,24 @@ def __mul__(self, other):
                                   unit=1.0 * result.units)
         return self.__class__(values=self._array * other, unit=self._unit)
 
+    def __matmul__(self, other):
+        if isinstance(other, self.__class__):
+            scale_l = self._unit.to_base_units()
+            scale_r = other._unit.to_base_units()
+            result = scale_l * scale_r
+            lhs = self._array
+            rhs = other._array * result.magnitude
+            lhs, rhs = self._broadcast(lhs, rhs)
+            return self.__class__(values=np.transpose(lhs @ rhs),
+                                  unit=1.0 * result.units)
+        if isinstance(other, Quantity):
+            scale_l = self._unit.to_base_units()
+            scale_r = other.to_base_units()
+            result = scale_l * scale_r
+            return self.__class__(values=self._array * result.magnitude,
+                                  unit=1.0 * result.units)
+        return self.__class__(values=self._array @ other, unit=self._unit)
+
     def __imul__(self, other):
         if isinstance(other, self.__class__):
             scale_l = self._unit.to_base_units()

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 cartesian_to_spherical, spherical_to_cartesian, rotation_matrix

src/osyris/spatial/spatial.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2022 Osyris contributors (https://github.com/nvaytet/osyris)
+
+import numpy as np
+from ..core import Array
+from ..config import units
+
+
+def rotation_matrix(vec, angle):
+    """
+    Returns 3D rotation matrix of angle 'angle' around rotation vector 'vec'.
+    """
+    if isinstance(vec, list):
+        vec = np.array(vec)
+    vec = vec / np.linalg.norm(vec)
+    R = np.cos(angle) * np.identity(3) + (np.sin(angle)) * np.cross(
+        vec,
+        np.identity(vec.shape[0]) * -1) + (1 - np.cos(angle)) * (np.outer(vec, vec))
+    return R
+
+
+def spherical_to_cartesian(r, theta, phi):
+    """
+    Converts spherical components radius, colatitude and azimuth to x,y,z
+    Returns an osyris array of shape (len(r), 3) containing a stacked xyz
+    """
+    x, y, z = (r * np.cos(phi) * np.sin(theta), r * np.sin(phi) * np.sin(theta),
+               r * np.cos(theta))
+    xyz = np.vstack([x, y, z])
+    return Array(values=np.transpose(xyz).values, unit=r.unit, name="position")
+
+
+def cartesian_to_spherical(position, origin=[0, 0, 0]):
+    """
+    Converts cartesian components in 'position' to spherical components
+    Returns a list of osyris Arrays containing radius, colatitude and azimuth
+    """
+    if isinstance(origin, Array):
+        centered_pos = position - origin
+    else:
+        centered_pos = position - Array(values=origin, unit=centered_pos.unit)
+
+    X = centered_pos[:, 0].values
+    Y = centered_pos[:, 1].values
+    Z = centered_pos[:, 2].values
+
+    radius = np.linalg.norm([X, Y, Z], axis=0)
+    colatitude = np.arctan2(np.linalg.norm([X, Y], axis=0), Z)
+
+    azimuth = np.arctan2(Y, X) + np.pi
+
+    radius = Array(values=radius, unit=position.unit, name='radius')
+    colatitude = Array(values=colatitude, unit=units('rad'), name='colatitude')
+    azimuth = Array(values=azimuth, unit=units('rad'), name='azimuth')
+
+    return radius, colatitude, azimuth
+
+
+def cartesian_to_cylindrical(position, origin=[0, 0, 0]):
+    """
+    Converts cartesian components in 'position' to cylindrical components
+    Returns a list of osyris Arrays containing radius, azimuth and elevation
+    """
+    if isinstance(origin, Array):
+        centered_pos = position - origin
+    else:
+        centered_pos = position - Array(values=origin, unit=centered_pos.unit)
+
+    elevation = centered_pos[:, 2].values
+
+    radius = np.linalg.norm([centered_pos[:, 0].values, centered_pos[:, 1].values],
+                            axis=0)
+    azimuth = np.arctan2(centered_pos[:, 1].values, centered_pos[:, 0].values)
+
+    radius = Array(values=radius, unit=centered_pos.unit, name='radius')
+    azimuth = Array(values=azimuth, unit=units('rad'), name='azimuth')
+    elevation = Array(values=elevation, unit=centered_pos.unit, name='elevation')
+
+    return radius, azimuth, elevation
+
+
+def cylindrical_to_cartesian(radius, azimuth, elevation):
+    """
+    Converts cylindrical components radius, azimuth and elevation to x,y,z
+    Returns an osyris array of shape (len(radius), 3) containing a stacked xyz
+    """
+    x, y, z = (radius * np.cos(azimuth), radius * np.sin(azimuth), elevation)
+    xyz = np.vstack([x, y, z])
+    return Array(values=np.transpose(xyz).values, unit=radius.unit, name="position")