Merge pull request #27 from jeonggyukim/read_hdf5

Minor enhancement to read hdf5

pyathena/io/read_hdf5.py

@@ -4,16 +4,19 @@
 
 import xarray as xr
 import numpy as np
+import h5py
 
 from .athena_read import athdf
 
-def read_hdf5(filename, **kwargs):
+def read_hdf5(filename, header_only=False, **kwargs):
     """Read Athena hdf5 file and convert it to xarray Dataset
 
     Parameters
     ----------
     filename : str
         Data filename
+    header_only : bool
+        Flag to read only attributes, not data.
     **kwargs : dict, optional
         Extra arguments passed to athdf. Refer to athdf documentation for
         a list of all possible arguments.
@@ -26,11 +29,29 @@ def read_hdf5(filename, **kwargs):
     See Also
     --------
     io.athena_read.athdf
+    load_sim.LoadSim.load_hdf5
+
+    Examples
+    --------
+    >>> from pyathena.io import read_hdf5
+    >>> ds = read_hdf5("/path/to/hdf/file")
+
+    >>> from pyathena.load_sim import LoadSim
+    >>> s = LoadSim("/path/to/basedir")
+    >>> ds = read_hdf5(s.files['hdf5']['prim'][30])
     """
+    if header_only:
+        with h5py.File(filename, 'r') as f:
+            data = {}
+            for key in f.attrs:
+                data[str(key)] = f.attrs[key]
+            return data
+
     ds = athdf(filename, **kwargs)
 
     # Convert to xarray object
-    varnames = set(map(lambda x: x.decode('ASCII'), ds['VariableNames']))
+    possibilities = set(map(lambda x: x.decode('ASCII'), ds['VariableNames']))
+    varnames = {var for var in possibilities if var in ds}
     variables = [(['z', 'y', 'x'], ds[varname]) for varname in varnames]
     attr_keys = (set(ds.keys()) - varnames
                  - {'VariableNames','x1f','x2f','x3f','x1v','x2v','x3v'})

pyathena/load_sim.py

@@ -275,6 +275,20 @@ def load_hdf5(self, num=None, ihdf5=None,
         Returns
         -------
         ds : xarray AthenaDataSet or yt datasets
+
+
+        Examples
+        --------
+        >>> from pyathena.load_sim import LoadSim
+        >>> s = LoadSim("/path/to/basedir")
+        >>> # Load everything at snapshot number 30.
+        >>> ds = s.load_hdf5(30)
+        >>> # Read the domain information only, without loading the fields.
+        >>> ds = s.load_hdf5(30, header_only=True)
+        >>> # Load the selected fields.
+        >>> ds = s.load_hdf5(30, quantities=['dens', 'mom1', 'mom2', 'mom3'])
+        >>> # Load the selected region.
+        >>> ds = s.load_hdf5(30, x1_min=-0.5, x1_max=0.5, x2_min=1, x2_max=1.2)
         """
 
         if num is None and ihdf5 is None:

pyathena/util/transform.py

@@ -1,8 +1,33 @@
 import numpy as np
 import xarray as xr
+from scipy.spatial.transform import Rotation
 
 
-def to_spherical(vec, origin):
+def euler_rotation(vec, angles):
+    """Rotate coordinate axes to transform the components of vector field `vec`
+
+    Parameters
+    ----------
+    vec : tuple-like
+        (vx, vy, vz) representing Cartesian vector components
+    angles : array
+        Euler angles [alpha, beta, gamma] in radian
+
+    Returns
+    -------
+    tuple
+        Cartesian components of the rotated vector.
+    """
+    angles = np.array(angles)
+    r = Rotation.from_euler('zyx', -angles, degrees=False)
+    rotmat = r.as_matrix()
+    vxp = rotmat[0, 0]*vec[0] + rotmat[0, 1]*vec[1] + rotmat[0, 2]*vec[2]
+    vyp = rotmat[1, 0]*vec[0] + rotmat[1, 1]*vec[1] + rotmat[1, 2]*vec[2]
+    vzp = rotmat[2, 0]*vec[0] + rotmat[2, 1]*vec[1] + rotmat[2, 2]*vec[2]
+    return (vxp, vyp, vzp)
+
+
+def to_spherical(vec, origin, newz=None):
     """Transform vector components from Cartesian to spherical coordinates
 
     Assumes vec is a tuple of xarray.DataArray.
@@ -13,6 +38,9 @@ def to_spherical(vec, origin):
         (vx, vy, vz) representing Cartesian vector components
     origin : tuple-like
         (x0, y0, z0) representing the origin of the spherical coords.
+    newz : array, optional
+        Cartesian components of the z-axis vector for the spherical
+        coordinates. If not given, it is assumed to be (0, 0, 1).
 
     Returns
     -------
@@ -24,22 +52,35 @@ def to_spherical(vec, origin):
     """
     vx, vy, vz = vec
     x0, y0, z0 = origin
-    x, y, z = vx.x, vx.y, vx.z
+    x, y, z = vx.x - x0, vx.y - y0, vx.z - z0
+
+    if newz is not None:
+        if ((np.array(newz)**2).sum() == 0):
+            raise ValueError("new z axis vector should not be a null vector")
+
+        # Rotate to align z axis
+        zhat = np.array([0, 0, 1])
+        alpha = np.arctan2(newz[1], newz[0])
+        beta = np.arccos(np.dot(newz, zhat) / np.sqrt(np.dot(newz, newz)))
+        vx, vy, vz = euler_rotation((vx, vy, vz), [alpha, beta, 0])
+        x, y, z = euler_rotation((x, y, z), [alpha, beta, 0])
 
     # Calculate spherical coordinates
-    R = np.sqrt((x-x0)**2 + (y-y0)**2)
-    r = np.sqrt(R**2 + (z-z0)**2)
-    th = np.arctan2(R, z-z0)
-    ph = np.arctan2(y-y0, x-x0)
+    R = np.sqrt(x**2 + y**2)
+    r = np.sqrt(R**2 + z**2)
+    th = np.arctan2(R, z)
+    ph = np.arctan2(y, x)
+
     # Move branch cut [-pi, pi] -> [0, 2pi]
     ph = ph.where(ph >= 0, other=ph + 2*np.pi)
-    sin_th, cos_th = R/r, (z-z0)/r
-    sin_ph, cos_ph = (y-y0)/R, (x-x0)/R
+    sin_th, cos_th = R/r, z/r
+    sin_ph, cos_ph = y/R, x/R
+
     # Avoid singularity
-    sin_th.loc[dict(x=x0, y=y0, z=z0)] = 0
-    cos_th.loc[dict(x=x0, y=y0, z=z0)] = 0
-    sin_ph.loc[dict(x=x0, y=y0)] = 0
-    cos_ph.loc[dict(x=x0, y=y0)] = 0
+    sin_th = sin_th.where(r != 0, other=0)
+    cos_th = cos_th.where(r != 0, other=0)
+    sin_ph = sin_ph.where(R != 0, other=0)
+    cos_ph = cos_ph.where(R != 0, other=0)
 
     # Transform Cartesian (vx, vy, vz) ->  spherical (v_r, v_th, v_phi)
     v_r = (vx*sin_th*cos_ph + vy*sin_th*sin_ph + vz*cos_th).rename('v_r')

pyathena/util/units.py

@@ -24,6 +24,7 @@ def __init__(self, kind='LV', muH=1.4271):
             self.length = 1.0
             self.mass = 1.0
             self.time = 1.0
+            self.units_override = None
             return
 
         mH = 1.008*au.u