Skip to content

Commit

Permalink
Merge pull request #27 from BayaneMdW/update
Browse files Browse the repository at this point in the history
add various utilities
  • Loading branch information
nicolasaunai authored May 2, 2024
2 parents 8567ac5 + 7cce194 commit 762687b
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 0 deletions.
4 changes: 4 additions & 0 deletions spok/coordinates/coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,3 +107,7 @@ def swi_base(vx, vy, vz, bx, by, bz):
X, Z = X.T, Z.T
Y = np.cross(Z, X)
return X, Y, Z

def rotates_from_phi_angle(x,y,z,angle):
r,th,ph = cartesian_to_spherical(x,y,z)
return spherical_to_cartesian(r,th,ph+angle)
58 changes: 58 additions & 0 deletions spok/utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import numpy as np
import pandas as pd
from scipy.interpolate import griddata
from scipy.ndimage import gaussian_filter


def listify(arg):
Expand Down Expand Up @@ -51,3 +53,59 @@ def make_center_bins(vec, dd = 1):
return [0.5*(v[1:,1:]+v[:-1,:-1]) for v in vec]
elif dd==3 :
return [0.5*(v[1:,1:,1:]+v[:-1,:-1,:-1]) for v in vec]


def reshape_to_2Darrays(lst_arrays):
lst_arrays = [np.array(listify(el)) for el in lst_arrays]
if np.sum([lst_arrays[0].shape!=el.shape for el in lst_arrays[1:]]):
raise ValueError('All elements of lst_arrays must have the same shape')
if len(lst_arrays[0].shape)==1 :
a2d = np.array(lst_arrays).T
old_shape = np.array(lst_arrays).shape
else :
a2d = np.asarray(lst_arrays)
old_shape = a2d.shape
a2d = a2d.T.ravel().reshape(np.prod(old_shape[1:]),old_shape[0])
return a2d, old_shape


def reshape_to_original_shape(a2d, old_shape):
if a2d.shape[0]==1 :
lst_arrays = a2d.T
else :
lst_arrays = np.asarray([a2d.reshape(old_shape[-1],np.prod(old_shape[:-1])).T[i::old_shape[0]] for i in range(old_shape[0])])
return lst_arrays


def regular_grid_interpolation(x, y, qty, new_x, new_y, **kwargs):
method = kwargs.get('method','linear')
qty_2d = reshape_to_2Darrays([qty])[0]
xy = reshape_to_2Darrays([x,y])[0]
reg_qty = griddata(xy, qty_2d[:,0], (new_x, new_y), method=method)
return reg_qty

def nan_gaussian_filter(arr, sigma, mode='nearest'):
"""Apply a gaussian filter to an array with nans.
Intensity is only shifted between not-nan pixels and is hence conserved.
The intensity redistribution with respect to each single point
is done by the weights of available pixels according
to a gaussian distribution.
All nans in arr, stay nans in gauss.
"""
nan_msk = np.isnan(arr)

loss = np.zeros(arr.shape)
loss[nan_msk] = 1
loss = gaussian_filter(
loss, sigma=sigma, mode=mode, cval=1)

gauss = arr / (1-loss)
gauss[nan_msk] = 0
gauss = gaussian_filter(
gauss, sigma=sigma, mode=mode, cval=0)
gauss[nan_msk] = np.nan

return gauss


0 comments on commit 762687b

Please sign in to comment.