Updating smileih5.py

CHANGELOG.md

@@ -5,6 +5,19 @@ current master
 --------------
 
 
+**Highlights**
+
+* Add support to read the `Smilei` PIC (https://smileipic.github.io/Smilei/) data format in both cartesian and azimuthal geometry. Postpic uses a build in azimuthal mode expansion very similar to the one used for fbpic.
+* To read smilei data, postpic only relies on the hdf5 package and not smilei's happi module for data access. Paricle ID's (ParticleTracking as described by smilei) can be read directly from the hdf5. Happi requires to sort the IDs and write a new hdf5, which can be twice as big as the original dumps. Using postpic's access this step will be skipped and thus access is much faster (but by default with unordered particle IDs as in any other code).
+
+
+**Incompatible adjustments to previous version**
+
+
+
+**Other improvements and new features**
+
+
 v0.5
 ----
 

postpic/datareader/smileih5.py

@@ -16,7 +16,7 @@
 #
 # Stephan Kuschel, 2023
 # Carolin Goll, 2023
-
+# Vinith Samson J, 2024
 
 from __future__ import absolute_import, division, print_function, unicode_literals
 
@@ -26,6 +26,10 @@
 import os.path
 import h5py
 import glob
+import numpy as np
+from .. import helper
+from ..helper_fft import fft
+import re
 
 __all__ = ['SmileiReader', 'SmileiSeries']
 
@@ -114,15 +118,174 @@ def __init__(self, h5file, iteration=None):
 
         if iteration is None:
             self._iteration = int(list(self._h5['data'].keys())[0])
+        elif iteration not in [int(i) for i in list(self._h5['data'].keys())]:
+            raise IOError("Iteration {} is in valid".format(iteration))
         else:
             self._iteration = int(iteration)
+
         self._data = self._h5['/data/{:010d}/'.format(self._iteration)]
         self.attrs = self._data.attrs
 
+    @staticmethod
+    def _modeexpansion_naiv(rawdata, theta=0):
+        '''
+        This method performes mode expansion of the raw data (an array consisting of complex
+        numbers) for both single and multiple theta vaues.
+
+        The output array has the shape (No.of theta, Nx, Nr)
+
+        Args:
+            rawdata : numpy array
+            The elements of this array are complex numbers, this got structured from the
+            raw data dumped in h5 file through getExpanded(key, theta) function.
+
+            theta : float/integer OR list of floats/integer
+
+        Output F_total is an array of real numbers which has shape (Np.of theta, Nx, Nr),
+        this F_total is the real value summation of the fourier series.
+        '''
+        if np.array(theta).shape is ():
+            theta = [theta]
+
+        array_list = []
+
+        (Nm, Nx, Nr) = rawdata.shape
+        F_total = np.zeros((Nx, Nr))
+        mode = [m for m in range(0, Nm)]
+        for t in theta:
+
+            for m in mode:
+                F_total += np.real(rawdata[m])*np.cos(m*t)+np.imag(rawdata[m])*np.sin(m*t)
+            array_list.append(F_total)
+
+        mod_F_total = np.stack(array_list, axis=0)
+        return mod_F_total
+
 # --- Level 0 methods ---
 
+    def _listAMmodes(self):
+        '''
+        This method is used to get the list of
+        [prefix of field names, No.of AM modes] available in the dump.
+        And it works only for AM mode technique.
+        '''
+        strings = np.array(self._data)
+        mask = np.array(["_mode_" in s for s in strings])
+        arr = strings[mask]
+        max_suffix = float('-inf')
+        max_suffix_string = None
+        prefix_list = []
+
+        for i in arr:
+            prefix, suffix = i.split('_mode_')
+            suffix_int = int(suffix)
+            if suffix_int > max_suffix:
+                max_suffix = suffix_int
+                max_suffix_string = i
+            prefix_list.append(prefix)
+
+        availModes = [i for i in range(0, int(max_suffix_string[-1])+1)]
+
+        # [field names prefix, available AM modes]
+        return [prefix_list, availModes]
+
 # --- Level 1 methods ---
 
+    def _getExpanded(self, key, theta=0):
+        '''
+        _getExpanded() method converts the raw data real number array from h5file into
+        a complex number array (This convertion is important while performing mode expansion)
+        and finally returns the mode expanded array.
+
+        This method takes input from the h5files dump which has the following format,
+        field array = [[real_1,imag_1,real_2,image_2,.....],...]
+        The shape of this field array is (Nx, 2x Nr)
+        After real->complex conversion, the field array shape takes the form (Nx, Nr)
+        This final array is fed into _modeexpansion_naiv method.
+        '''
+        array_list = []
+        modes = self._listAMmodes()[-1]
+        for mode in modes:
+
+            field_name = key+"_mode_"+str(mode)
+            field_array = np.array(self._data[field_name])
+            field_array_shape = field_array.shape
+            reshaped_array = field_array.reshape(field_array_shape[0], field_array_shape[1]//2, 2)
+            complex_array = reshaped_array[:, :, 0] + 1j * reshaped_array[:, :, 1]
+            array_list.append(complex_array)
+
+        # Modified array of shape (Nmodes, Nx, Nr)
+        mod_complex_data = np.stack(array_list, axis=0)
+        factor = self._data["{}_mode_0".format(key)].attrs['unitSI']
+        return SmileiReader._modeexpansion_naiv(rawdata=mod_complex_data, theta=theta)*factor
+
+    def data(self, key, **kwargs):
+        '''
+        should work with any key, that contains data, thus on every hdf5.Dataset,
+        but not on hdf5.Group. Will extract the data, convert it to SI and return it
+        as a numpy array. Constant records will be detected and converted to
+        a numpy array containing a single value only.
+
+        If the key is in AM mode dump, then it performs the mode expansion.
+        The theta values for which we need to perform
+        mode expansion can be given as keyword args.
+
+        Example:
+            Data = data(key='El', theta=[0,pi/2,pi])
+            Now the Data will array have the shape (3, Nx, Nr)
+        '''
+
+        # checking whether the key is in AM mode dump
+        if key in self._listAMmodes()[0]:
+            return self._getExpanded(key=key, **kwargs)
+        else:
+            record = self._data[key]
+
+        if "value" in record.attrs:
+            # constant data (a single int or float)
+            ret = np.float64(record.attrs['value']) * record.attrs['unitSI']
+        else:
+            # array data
+            ret = np.float64(record[()]) * record.attrs['unitSI']
+        return ret
+
+    # To get the offsets of the grid.
+    def gridoffset(self, key, axis):
+        axid = helper.axesidentify[axis]
+
+        if axid == 91:  # theta
+            return 0
+        elif key in self._listAMmodes()[0] and axid in [0, 90]:
+            key = "{}_mode_0".format(key)
+            axid = int(axid/90)
+            return super(SmileiReader, self).gridoffset(key=key, axis=axid)
+        else:
+            return super(SmileiReader, self).gridoffset(key, axis)
+
+    # To get the grid spacing.
+    def gridspacing(self, key, axis):
+        axid = helper.axesidentify[axis]
+
+        if key in self._listAMmodes()[0] and axid in [0, 90]:
+            key = "{}_mode_0".format(key)
+            axid = int(axid/90)
+            return super(SmileiReader, self).gridspacing(key=key, axis=axid)
+        else:
+            return super(SmileiReader, self).gridspacing(key, axis)
+
+    # To get the grid points
+    def gridpoints(self, key, axis):
+        axid = helper.axesidentify[axis]
+
+        if key in self._listAMmodes()[0]:
+            key = "{}_mode_0".format(key)
+            (Nx, Nr) = self._data[key].shape
+            Nr = Nr/2
+            axid = int(axid/90)
+            return (Nx, Nr)[axid]
+        else:
+            return super(SmileiReader, self).gridpoints(key=key, axis=axis)
+
 # --- Level 2 methods ---
 
     def _keyE(self, component, **kwargs):
@@ -137,8 +300,9 @@ def _keyB(self, component, **kwargs):
 
     def _simgridkeys(self):
         # Smilei deviates from openPMD standard: Ex instead of E/x
-        return ['Ex', 'Ey', 'Ez',
-                'Bx', 'By', 'Bz']
+        return ['Ex', 'Ey', 'Ez', 'Er', 'El', 'Et',
+                'Bx', 'By', 'Bz', 'Br', 'Bl', 'Bt',
+                'Jx', 'Jy', 'Jz', 'Jr', 'Jl', 'Jt', 'Rho']
 
     def getderived(self):
         '''