From 7cce19497033f718fdb46c830e2878374cd1e341 Mon Sep 17 00:00:00 2001
From: BayaneMdW <bayane.michotte-de-welle@lpp.polytechnique.fr>
Date: Thu, 2 May 2024 13:28:30 +0000
Subject: [PATCH] add various utilities

---
 spok/coordinates/coordinates.py |  4 +++
 spok/utils.py                   | 58 +++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/spok/coordinates/coordinates.py b/spok/coordinates/coordinates.py
index 35e9a7e..6872e02 100644
--- a/spok/coordinates/coordinates.py
+++ b/spok/coordinates/coordinates.py
@@ -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)
diff --git a/spok/utils.py b/spok/utils.py
index c0a84e0..3f4266f 100644
--- a/spok/utils.py
+++ b/spok/utils.py
@@ -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):
@@ -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
+
+