diff --git a/epix/electric_field_handler.py b/epix/electric_field_handler.py index 95228ee..5b669fe 100755 --- a/epix/electric_field_handler.py +++ b/epix/electric_field_handler.py @@ -30,28 +30,32 @@ def __init__(self, field_map=""): self._build_interpolator() else: raise ValueError(f'Cannot open "{self.map}". It is not a valid file' - ' for the electirc field map.') + ' for the electric field map.') def _load_field(self): file_ending = self.map.split('.')[-1] if file_ending == 'csv': - self.field = pd.read_csv(self.map) + _field = pd.read_csv(self.map) + _field = pd.DataFrame(_field.groupby(['r']).aggregate({'z': list, 'E': list})) + _field = _field.explode(['z', 'E']) + _field = _field.reset_index() + self.field = _field.applymap(float) elif file_ending == 'gz': with gzip.open(self.map, 'rb') as f: field_map = json.load(f) - + csys = field_map['coordinate_system'] grid = [np.linspace(left, right, points) - for _,(left, right, points) in csys] + for _, (left, right, points) in csys] csys = np.array(np.meshgrid(*grid, indexing='ij')) axes = np.roll(np.arange(len(grid) + 1), -1) csys = np.transpose(csys, axes) csys = np.array(csys).reshape((-1, len(grid))) - - self.field = pd.DataFrame() - self.field["r"] = np.array(csys)[:,0] - self.field["z"] = np.array(csys)[:,1] + + self.field = pd.DataFrame() + self.field["r"] = np.array(csys)[:, 0] + self.field["z"] = np.array(csys)[:, 1] self.field["E"] = np.array(field_map['map']) else: raise ValueError(f'Cannot open "{self.map}". File extension is not valid' @@ -63,12 +67,12 @@ def _get_coordinates(self): def _build_interpolator(self): e_tmp = np.reshape(np.array(self.field.E), - (len(self.Z), len(self.R))) - self.interpolator = RGI([self.Z, self.R], + (len(self.R), len(self.Z))) + self.interpolator = RGI([self.R, self.Z], e_tmp, bounds_error=False, fill_value=None) - + def get_field(self, x, y, z, outside_map=np.nan): """ Function which returns the electric field at a certain position @@ -84,7 +88,7 @@ def get_field(self, x, y, z, outside_map=np.nan): was not within the range of the map. Default np.nan :return: """ - r = np.sqrt(x**2+y**2) - efield = self.interpolator((z, r)) + r = np.sqrt(x ** 2 + y ** 2) + efield = self.interpolator((r, z)) efield[np.isnan(efield)] = outside_map return efield