barsukov 1.3.3__py3-none-any.whl → 1.3.5__py3-none-any.whl

barsukov/__init__.py

@@ -1,18 +1,15 @@
 # Modules:
 from . import time
 from . import data
-    
+
 
 # Objects/Functions:
 from .script import Script
 from .logger import Logger
-
 from .obj2file import *
 
 
 # Equipment Objects:
 from .exp.mwHP import mwHP
 
-__all__ = ["time", "data", "save_object", "load_object", "Script", "Logger", "mwHP"]
-
 

barsukov/data/Change_phase.py

@@ -0,0 +1,160 @@
+from barsukov.data.constants import deg2rad
+from barsukov.time import time_stamp
+
+from scipy.optimize import curve_fit
+from scipy.optimize import differential_evolution
+import numpy as np
+import matplotlib.pyplot as plt
+
+import glob
+import sys
+import os
+
+
+class Change_phase:
+        ### the phase you receive from auto, is a phase shift that you need to add to the phase of the original data. 
+        ### adding the two phases together gives to you the total effective lock-in phase of the calculated data.
+        ### Note that lock-in phase corresponds to the reference, not to the signal itself.
+        ### Lock-in phase is an artificial phase delay of the reference
+        ### new reference is cos(Wt - phase)
+        ### This script's phase, if added to the original phase of the lock-in, will give you a cumulative phase.
+        ### The recalculated signal would correspond to lock-in signal if measured with this cumulative phase.
+        ### This cumulative phase is the phase delay of your signal with respect to the original unaltered reference.
+        ### The automatically recalculated data is correct only if considered together with the automatically calculated phase
+        ### This means, you may get positive or negative signals in x-channel. So always consider the cummulative phase when evaluating the data. 
+
+    def __init__(self, x=[], A=[], B=[], initial_phase=0):
+        self.x = np.array(x)
+        self.A = np.array(A)
+        self.B = np.array(B)
+        self.initial_phi = initial_phase
+
+        self.phi = initial_phase
+        self.newA = None
+        self.newB = None
+
+
+    def read_from_file(self, full_file_path, x_column=0, A_column=1, B_column=2, initial_phase=0):
+        self.full_file_path = full_file_path
+        data = np.loadtxt(self.full_file_path, skiprows=0, unpack=True, usecols=(x_column, A_column, B_column))
+        self.x = data[0]
+        self.A = data[1]
+        self.B = data[2]
+        self.initial_phi = initial_phase
+
+
+    def offset_phase(self, phi=None):
+        if phi is not None:
+            self.phi = float(phi)
+            self.newA = self.A * np.cos(self.phi*deg2rad) + self.B * np.sin(self.phi*deg2rad)
+            self.newB = - self.A * np.sin(self.phi*deg2rad) + self.B * np.cos(self.phi*deg2rad)
+        else:
+            def to_minimize(phi_val):
+                self.phi = phi_val[0] #Differential evolution passes arrays
+                self.newA = self.A * np.cos(self.phi*deg2rad) + self.B * np.sin(self.phi*deg2rad)
+                self.newB = - self.A * np.sin(self.phi*deg2rad) + self.B * np.cos(self.phi*deg2rad)
+                popt, pcov = curve_fit(lambda x,a,b: a+b*x, self.x, self.newB, p0=[0,0])
+                return 1-(pcov[0,1]/(pcov[0,0]*pcov[1,1]))**2
+
+            ### MINIMIZES the sum of the data in the Y-channel (B)
+            result = differential_evolution(to_minimize, bounds=[(0,359.99)], strategy='best1bin')
+            if result.success:
+                self.offset_phase(phi=result.x[0])
+                print(f"Auto-adjusted phase: {result.x[0]} degrees")
+            else:
+                self.phi = initial_phase
+                print("Optimization failed!")
+
+
+    def plot_offset(self):
+        if not hasattr(self, "fig") or not hasattr(self, "axes"):
+            self.fig, self.axes = plt.subplots(nrows=4, ncols=1, figsize=(12,18))
+
+            self.lines = [
+                self.axes[0].plot(self.x, self.A ,'b-', label='X-')[0],
+                self.axes[0].plot(self.x, self.B, 'r-', label='Y-')[0],
+                self.axes[1].plot(self.x, self.newA, 'b-', label='X-')[0],
+                self.axes[2].plot(self.x, self.newB, 'r-', label='Y-')[0],
+                self.axes[3].plot(self.x, self.newA ,'b-', label='X-')[0],
+                self.axes[3].plot(self.x, self.newB, 'r-', label='Y-')[0],
+            ]
+
+            self.axes[0].set_title('Original X- & Y-')
+            self.axes[1].set_title('Adjusted X-')
+            self.axes[2].set_title('Adjusted Y-')
+            self.axes[3].set_title('Adjusted X- & Y-')
+
+            for ax in self.axes:
+                ax.legend()
+        else:
+            for line in self.lines:
+                line.set_xdata(self.x)
+
+            self.lines[0].set_ydata(self.A)
+            self.lines[1].set_ydata(self.B)
+            self.lines[2].set_ydata(self.newA)
+            self.lines[3].set_ydata(self.newB)
+            self.lines[4].set_ydata(self.newA)
+            self.lines[5].set_ydata(self.newB)
+
+            for ax in self.axes:
+                ax.relim()
+                ax.autoscale_view()
+
+        if "IPython" in sys.modules:
+            from IPython.display import display
+            display(self.fig)
+        else:
+            plt.ion()
+            self.fig.canvas.draw()
+            self.fig.canvas.flush_events()
+            plt.show()
+
+
+    def save_data(self, full_folder_path=None, file_name=None):
+        if hasattr(self, "full_file_path"):
+            full_folder_path, file_name = os.path.split(self.full_file_path)
+            file_name = f"Corrected_{round(self.phi)}_{file_name}"
+        else:
+            if full_folder_path is None:
+                full_folder_path = os.getcwd()
+            if file_name is None:
+                file_name = f"{time_stamp()}_Corrected_Phase_Lock_in_Data"
+
+        full_folder_path = os.path.join(full_folder_path, 'phase-corrected_data')
+        if not os.path.isdir(full_folder_path):
+            os.makedirs(full_folder_path)
+
+        full_file_path = os.path.join(full_folder_path, file_name)
+        with open(full_file_path, "w") as file:
+            for i in range(len(self.x)):
+                file.write(f"{self.x[i]} {self.newA[i]} {self.newB[i]} \n")
+
+
+    def offset_phase_script(self, full_folder_path=None, x_column=0, A_column=1, B_column=2, initial_phase=0, nocheck=False):
+        if full_folder_path is None:
+            full_folder_path = os.getcwd()
+        for file in glob.glob(os.path.join(full_folder_path, '*.txt')):
+            self.read_from_file(file, x_column, A_column, B_column, initial_phase)
+
+            self.offset_phase()
+            self.plot_offset()
+
+            while True:
+                if nocheck is False:
+                    manual_input=input('Enter "auto" to auto-calculate phase, a number/float for manual phase, or press ENTER to skip: ')
+
+                    if manual_input == "auto":
+                        self.offset_phase()
+                        self.plot_offset()
+                    elif manual_input:
+                        try:
+                            self.offset_phase(manual_input)
+                            self.plot_offset()
+                        except ValueError:
+                            print(f"Invalid input: {manual_input}. Please enter a valid phase value or 'auto'.")
+                    if not manual_input:
+                        print("Phase adjustment completed.")
+                        break
+            self.save_data()
+            print("Adjusted phase data saved.")

barsukov/data/Lock_in_emulator.py

@@ -0,0 +1,175 @@
+import numpy as np
+from barsukov.data import noise
+
+import sympy as sp
+
+class Lock_in_emulator:
+    def __init__(self, signal, f, phase, x_start, x_stop, x_amp, time, dt, TC, order, plot_points, buffer_size, 
+                 johnson_T=0, johnson_R=0,
+                 shot_I=0, shot_R=0,
+                 onef_rms=0,
+                 rtn_tau_up=0, rtn_tau_down=0, rtn_state_up=0, rtn_state_down=0,
+                 bit_depth=0, bit_measure_min=0, bit_measure_max=0):
+
+        #Signal Properties
+        self.signal_arr = signal
+        self.f = f
+        self.phase = np.pi * phase / 180
+        self.x_start = x_start
+        self.x_stop = x_stop
+        self.x_amp = x_amp
+        self.time = abs(time)
+
+        #Noise Properties
+        self.jT, self.jR = johnson_T, johnson_R
+        self.sI, self.sR = shot_I, shot_R
+        self.oRMS = onef_rms
+        self.rTU, self.rTD, self.rSU, self.rSD = rtn_tau_up, rtn_tau_down, rtn_state_up, rtn_state_down
+        self.bD, self.bMIN, self.bMAX = bit_depth, bit_measure_min, bit_measure_max
+
+        #Filter Properties
+        self.dt = abs(dt)
+        self.TC = TC
+        self.n = abs(order)
+
+
+        #Plotting Properties
+        self.plot_points = plot_points
+        self.buffer_size = buffer_size
+        self.buffer_period = self.buffer_size*self.dt
+        self.buffer_offset = np.linspace(-self.buffer_period, 0, self.buffer_size)
+
+    def run(self):
+        self.t_plot = np.linspace(self.buffer_period, self.time, self.plot_points)
+        self.x_plot = self.x_arr(self.t_plot)
+
+        self.original_signal = self.signal_arr(self.x_plot)
+        self.expected_signal = 0.5 * self.x_amp * np.gradient(self.original_signal, self.x_plot)
+        self.output_signal = self.signal_output_arr(self.t_plot)
+
+        self.fit()
+        #self.plot()
+
+### BEGIN: FIELD
+    def x_arr(self, t_arr):
+        return self.x_start + t_arr * (self.x_stop - self.x_start) / self.time
+
+    def x_with_mod_arr(self, t_arr):
+        return self.x_arr(t_arr) + self.x_amp * np.cos(2 * np.pi * self.f * t_arr)
+
+
+### BEGIN: NOISE
+    def noise(self, t_arr):
+        if self.jT and self.jR: self.johnson = noise.johnson(t_arr, self.jT, self.jR)
+        else: self.johnson = np.zeros(len(t_arr))
+        
+        if self.sI and self.sR: self.shot = noise.shot(t_arr, self.sI, self.sR)
+        else: self.shot = np.zeros(len(t_arr))
+            
+        if self.oRMS: self.onef = noise.color(t_arr, self.oRMS)
+        else: self.onef = np.zeros(len(t_arr))
+        
+        if (self.rTU and self.rTD) or (self.rSU and self.rSD): self.rtn = noise.rtn(t_arr, self.rTU, self.rTD, self.rSU, self.rSD)
+        else: self.rtn = np.zeros(len(t_arr))
+        return self.johnson + self.shot + self.onef + self.rtn
+
+
+    def lp_filter_arr(self, t_arr):
+        t_arr = -(t_arr - t_arr[-1])
+        factorial = np.math.factorial(self.n - 1)
+        lp_filter_arr = (t_arr ** (self.n - 1)) * np.exp(-t_arr / self.TC) / (self.TC**self.n * factorial)
+        return lp_filter_arr / abs(np.sum(lp_filter_arr) * self.dt)
+
+    def signal_output(self, t):
+        t_arr = t + self.buffer_offset
+
+        x_arr = self.x_with_mod_arr(t_arr)
+        s_arr = self.signal_arr(x_arr)
+        noise_arr = self.noise(t_arr)
+        filter_arr = self.lp_filter_arr(t_arr)
+        ref_X = np.cos(2 * np.pi * self.f * t_arr - self.phase)
+
+        integrand = (s_arr+noise_arr) * ref_X * filter_arr * self.dt
+        return np.sum(integrand)
+
+    def signal_output_arr(self, t_arr):
+        output = np.array([self.signal_output(t) for t in t_arr])
+
+        if self.bD or (self.bMIN and self.bMAX):
+            return noise.bit(output, self.bD, self.bMIN, self.bMAX)
+        else:
+            return output
+
+
+### BEGIN: FIT:
+    def fit(self):
+        from scipy.interpolate import interp1d
+        from scipy.optimize import curve_fit
+
+        interp = interp1d(self.x_plot, self.expected_signal, kind='cubic', fill_value=0.0, bounds_error=False)
+
+        def model(x, diminish, stretch, shift):
+            x_trans = (x / stretch) - shift
+            return (1.0 / diminish) * interp(x_trans)
+
+        expected_max_idx, expected_min_idx = np.argmax(self.expected_signal), np.argmin(self.expected_signal)
+        output_max_idx, output_min_idx = np.argmax(self.output_signal), np.argmin(self.output_signal)
+
+        expected_peak_xdif = np.abs(self.x_plot[expected_max_idx] - self.x_plot[expected_min_idx])
+        output_peak_xdif = np.abs(self.x_plot[output_max_idx] - self.x_plot[output_min_idx])
+
+        initial_stretch = np.abs(output_peak_xdif / expected_peak_xdif)
+        initial_diminish = np.ptp(self.expected_signal) / np.ptp(self.output_signal)
+
+        initial_transform = model(self.x_plot, initial_diminish, initial_stretch, 0)
+        transform_max_idx = np.argmax(initial_transform)
+
+        initial_shift = (self.x_plot[output_max_idx] - self.x_plot[transform_max_idx]) / initial_stretch
+
+        initial_guess = [np.max([1,initial_diminish]), np.max([1,initial_stretch]), np.max([-np.ptp(self.x_plot),initial_shift])]
+
+        bounds = [(1, 1, -np.ptp(self.x_plot)), (10, 0.5*np.ptp(self.x_plot), np.ptp(self.x_plot) )]
+        popt, _ = curve_fit(model, self.x_plot, self.output_signal, sigma=1e-4, p0=initial_guess, bounds=bounds, method='trf')
+
+        #print(f"initial diminish: {initial_diminish}\ninitial stretch: {initial_stretch}\ninitial shift: {initial_shift}")
+        self.diminish = popt[0]
+        self.stretch = popt[1]
+        self.shift = popt[2]
+        self.adjusted_signal = model(self.x_plot, self.diminish, self.stretch, self.shift)
+
+        #SNR Calculation
+        p2p = np.max(self.adjusted_signal) - np.min(self.adjusted_signal)
+        noise_sample = self.output_signal[int(len(self.adjusted_signal) * 0.9):] #last 10% of the output signal
+        v_rms = np.sqrt(np.mean(noise_sample**2))
+
+        self.snr = p2p / v_rms
+
+    def plot(self):
+        import matplotlib.pyplot as plt
+
+        print(f'diminish: {self.diminish}')
+        print(f'stretch: {self.stretch}')
+        print(f'shift: {self.shift}')
+        print(f'Signal to Noise Ratio: {self.snr}')
+
+        self.fig, self.axes = plt.subplots(nrows=2, ncols=1, figsize=(12,18))
+
+        self.lines = [
+            self.axes[0].plot(self.x_plot, self.original_signal, 'r-', label='Original Signal')[0],
+            self.axes[1].plot(self.x_plot, self.output_signal, 'b-', label='Demodulated Signal (Lock-In)')[0],
+            self.axes[1].plot(self.x_plot, self.expected_signal, 'r-', label='Demodulated Signal (Expected)')[0],
+            self.axes[1].plot(self.x_plot, self.adjusted_signal, 'g-', label=f'Demodulated Signal (Adjusted)\n  Diminish: {self.diminish}\n  Stretch:{self.stretch}\n  Shift: {self.shift}')[0],
+        ]
+
+        self.axes[0].set_title('Original Signal vs x')
+        self.axes[1].set_title('Demodulated Signal vs x')
+
+        plt.legend()
+        plt.show()
+
+    def spectrum_average(self, num):
+        result = np.zeros_like(self.t_plot)
+        for i in range(0, num):
+            result += self.signal_output_arr(self.t_plot)
+
+        return result / num

barsukov/data/__init__.py

@@ -1 +1,6 @@
+from . import constants
+from . import noise
 from .fft import *
+from .Change_phase import Change_phase
+from .Lock_in_emulator import Lock_in_emulator
+from . import lock_in_emulator_app

barsukov/data/constants.py

@@ -0,0 +1,10 @@
+import numpy as np
+
+
+### Conversions:
+deg2rad=2*np.pi/360.0
+
+
+### Constants:
+K_b = 1.380649e-23 #Boltzmann constant
+q = 1.602176634e-19 # Elementary Charge

barsukov/data/fft.py

@@ -1,68 +1,114 @@
-### BEGIN Dependencies ###
 import numpy as np
-import scipy as sp
-from barsukov.logger import debug    # Un-comment before implementing in pip
-### END Dependencies ###
 
 
-def fft(x, y, equidistant_check=True, equidistant_rel_error=1e-4, remove_negative_f=False, inverse=False):
-    ### Takes: x=list of real floats, y=list of data or list of lists of data. Data can be real or complex.
-    ### Returns: freqs, fft_y_norm, msg
-    ### This fft used to give wrong sign of the imaginary component because of the "wrong" definition of fft in np and sp
-    ### Now it has been corrected to the Mathematica definition of fft = int ... exp(2pi f t)
-    msg = ''
+def fft(x, y, equidistant_check=True, equidistant_rel_error=1e-4, remove_negative_f=False, mathematica_convention=False, inverse=False):
+    """
+    Perform Fast Fourier Transform (FFT) or Inverse FFT on one-dimensional or multi-dimensional data,
+    with optional handling for non-equidistant sampling, mandatory normalization to continuous Fourier definition.
+    Allows for +- in exp for scientific and engineering definition for Fourier Transform.
+
+    Parameters
+    ----------
+    x (array_like): 
+        The time or spatial domain axis (1D array). Must be increasing, doesn't need to be equidistant.
+    y (array_like): 
+        The signal to be transformed. Can be 1D or 2D array (if 2D, the FFT is applied along axis 1).
+    equidistant_check (bool, optional): 
+        If True (default), checks whether `x` is uniformly spaced. If not, interpolates to uniform spacing.
+    equidistant_rel_error : float, optional
+        Relative tolerance for equidistant spacing check. Default is 1e-4.
+    remove_negative_f : bool, optional
+        If True, removes negative frequencies from output. Default is False.
+    mathematica_convention : bool, optional
+        If True, uses Mathematica-style conventions: forward Fourier ~ exp(i2pift).
+        If False (default), uses NumPy convention: forward Fourier ~ exp(-i2pift)
+    inverse : bool, optional
+        If True, computes the inverse FFT. Default is False (forward transform).
+
+    Returns
+    -------
+    fft_x : ndarray
+        Ordered frequency domain axis corresponding to the transformed data (negative frequencies can be removed if remove_negative_f is True).
+        Ordered time domain axis if inverse is True (time axis starts at 0).
+    fft_y : ndarray
+        Transformed signal. Shape matches `y`.
+
+    Notes
+    -----
+    - The function supports both real and complex-valued input `y`.
+    - If `x` is not equidistant and `equidistant_check=True`, the function interpolates `y` using `make_equidistant`.
+
+    Examples
+    --------
+    >>> t = np.linspace(0, 1, 1000)
+    >>> y = np.sin(2 * np.pi * 50 * t)
+    >>> f, Y = fft(t, y)
+
+    >>> # Inverse transform
+    >>> t_rec, y_rec = fft(f, Y, inverse=True)
+    """
+
+    x, y = np.array(x), np.array(y)
+
+    # Handle non-equidistant input
     if equidistant_check:
         diffs = np.diff(x)
         if not np.allclose(diffs, diffs[0], rtol=equidistant_rel_error):
             # x is not equidistant, must start interpolating
-            x,y = make_equidistant(x, y, step=None)
-            y = y.T
-            msg += debug('fft(x,y) made x,y equidistant.')
-
-    y = np.array(y)
-    #print(y)
+            x, y = make_equidistant(x, y, step=None)
 
+    # Determine shape and axis
     if y.ndim == 1:
-        # y is 1D, treat it as a single column
-        n = len(y)  # Number of points in the column
-        if inverse is False: fft_y = np.fft.ifft(y) * n # np fft has the "wrong" imag sign
-        else: fft_y = np.fft.fft(y) # That's why fft and ifft are inverted in this code
+        n, axis = len(y), -1
+    else:
+        n, axis = y.shape[1], 1
+
+    #Determine Convention
+    if mathematica_convention is False:
+        fft_func, fft_norm = np.fft.fft, "backward"
+        ifft_func, ifft_norm = np.fft.ifft, "forward"
     else:
-        # y is 2D, treat it as multiple columns
-        n = y.shape[1]  # Number of points in each column
-        if inverse is False: fft_y = np.fft.ifft(y, axis=1) * n # np fft has the "wrong" imag sign
-        else: fft_y = np.fft.fft(y, axis=1) # That's why fft and ifft are inverted in this code
-
-    sample_spacing = ( x[-1] - x[0] ) / (n-1)
-    #print(n, sample_spacing, x[1] - x[0])
-    fft_y_norm = fft_y * sample_spacing # This normalizes FFT to mathematically correct
-    freqs = np.fft.fftfreq(n, d=sample_spacing)
-
-    if remove_negative_f is False:
-        sorted_indices = np.argsort(freqs)
-        freqs = freqs[sorted_indices]
-        if isinstance(fft_y_norm[0], (list, np.ndarray)):  # If fft_y_norm contains multiple columns
-            fft_y_norm = [x[sorted_indices] for x in fft_y_norm]
-        else:  # If fft_y_norm is a single column
-            fft_y_norm = fft_y_norm[sorted_indices]
-        msg += debug('fft(x,y) sorted negative and positive frequencies.')
+        fft_func, fft_norm = np.fft.ifft, "forward"
+        ifft_func, ifft_norm = np.fft.fft, "backward"
+
+    # Compute FFT and normalize
+    sample_spacing = (x[-1] - x[0]) / (n-1.0)
+
+    if inverse is True:
+        y = np.fft.ifftshift(y, axes=axis) # Reorder y to match np.fft.ifft convention
+        fft_x = np.arange(0, n) / (n * sample_spacing) # Time domain creation
+        fft_y = ifft_func(y, axis=axis, norm=ifft_norm) * sample_spacing # Sample Spacing is the Fourier consistent normalization
     else:
-        mask = freqs >= 0 # Boolean array with Falses for negative frequencies, effectively removing them
-        freqs = freqs[mask]
-        # If fft_y_norm contains multiple columns:
-        if isinstance(fft_y_norm[0], (list, np.ndarray)):
-            fft_y_norm = [x[mask] for x in fft_y_norm]
-        else:  # If fft_y_norm is a single column
-            fft_y_norm = fft_y_norm[mask]
-        msg += debug('fft(x,y) removed negative frequencies.')
+        fft_x = np.fft.fftfreq(n, d=sample_spacing) # Frequency domain creation
+        fft_y = fft_func(y, axis=axis, norm=fft_norm) * sample_spacing
+        fft_x, fft_y = np.fft.fftshift(fft_x), np.fft.fftshift(fft_y, axes=axis) # Reorder x, y for plotting
+
+    # Remove negative frequencies
+    if remove_negative_f is True:
+        mask = fft_x >= 0
+        fft_x = fft_x[mask]
+        if y.ndim == 1:
+            fft_y = fft_y[mask]
+        else:
+            fft_y = fft_y[:, mask]
+
+    return fft_x, fft_y
+
+# NOTES For Developer:
+        # Mathematica has convention exp(i*2pi*t)
+        # np.fft.fft takes x input -3,-2,-1,0,1,2,3 and np.fft.ifft takes x input 0,1,2,3,-3,-2,-1
+        # Our fft(np.fft.ifft) takes x input -3,-2,-1,0,1,2,3 and our ifft(np.fft.fft) takes x input 0,1,2,3,-3,-2,-1
+        # Numpy outputs in 0,1,2,3,-3,-2,-1 order, our function sorts to -3,-2,-1,0,1,2,3
+
+        #Tested with Lorentzian, Gaussian, and Rectangular Signals
+
+def ifft(x, y, equidistant_check=True, equidistant_rel_error=1e-4, remove_negative_f=False, mathematica_convention=False):
+    return fft(x, y, equidistant_check=equidistant_check, equidistant_rel_error=equidistant_rel_error, remove_negative_f=remove_negative_f, mathematica_convention=mathematica_convention, inverse=True)
 
-    msg += debug('freqs, fft_y_norm, msg = fft(x,y) is done.\nThe forward fft approximates the mathematically correct integral over ...exp(+i2pift).\nNow do not forget to apply np.abs(fft_y_norm), np.angle(fft_y_norm), fft_y_norm.real, fft_y_norm.imag')
-    return freqs, fft_y_norm, msg
 
-def ifft(x, y, equidistant_check=True, equidistant_rel_error=1e-4, remove_negative_f=False):
-    return fft(x, y, equidistant_check=equidistant_check, equidistant_rel_error=equidistant_rel_error, remove_negative_f=remove_negative_f, inverse=True)
 
 def make_equidistant(x, y, step=None):
+    import scipy.interpolate as sp
     ### Takes one column x and one or more columns y and makes them equidistant in x
     ### Returns new_x, new_y. The number of points will likely change.
     if step is None:
@@ -70,18 +116,17 @@ def make_equidistant(x, y, step=None):
         min_step = np.min(np.diff(x))
     else:
         min_step = step
-    
+
     # Generate the new equidistant x array
     new_x = np.arange(x[0], x[-1] + min_step, min_step)
-    
+
     if isinstance(y[0], (list, np.ndarray)):  # If y contains multiple columns
         new_y = []
         for y_column in y:
             interpolation_function = sp.interpolate.interp1d(x, y_column, kind='linear', fill_value='extrapolate')
             new_y.append(interpolation_function(new_x))
-        new_y = np.array(new_y).T  # Transpose to match the original structure
     else:  # If y is a single column
         interpolation_function = sp.interpolate.interp1d(x, y, kind='linear', fill_value='extrapolate')
         new_y = interpolation_function(new_x)
-    
-    return new_x, new_y
+
+    return np.array(new_x), np.array(new_y)