Repository: TomHilder/wakeflow (language: python)

Issue:
Small suggestions on draft paper



Discussion / hints:

Fix patch:
diff --git a/src/wakeflow/burgers.py b/src/wakeflow/burgers.py
index 0439ffe..d8748d6 100644
--- a/src/wakeflow/burgers.py
+++ b/src/wakeflow/burgers.py
@@ -7,6 +7,8 @@
 
 import numpy                as np
 import matplotlib.pyplot    as plt
+import matplotlib.colors    as cl
+import matplotlib.cm        as cm
 
 # NOTE: contents are intended for internal use and should not be directly accessed by users
 
@@ -23,14 +25,15 @@ def _solve_burgers(
     linear_solution, 
     linear_t, 
     show_teta, 
-    tf_fac
+    tf_fac,
+    t_edge
 ): 
     """Propagate the wake in (t,eta,chi) space by solving Eq. 10 from Bollati et al. 2021 using Godunov scheme.
     """
 
-    # Eq. (15) Bollati et al. 2021
-    profile = profile * (gamma + 1) * beta_p / 2**(3 / 4)
-    deta    = eta[1] - eta[0]
+    #Get eta step
+    deta = eta[1] - eta[0]
+    #print(len(eta))
 
     # time required to develop N-wave for betap = 1
     tf_th = 300 
@@ -38,7 +41,8 @@ def _solve_burgers(
     # time required to display N-wave for generic betap, Eq. (39) Rafikov 2002
     tf  = tf_th / beta_p 
     tf *= tf_fac
-
+    
+    """
     # Eq. (18) Bollati et al. 2021
     eta_min = -eta_tilde - np.sqrt(2 * C * tf) - 3 
     eta_max = -eta_tilde + np.sqrt(2 * C * tf) + 3
@@ -65,7 +69,7 @@ def _solve_burgers(
 
         #if show_teta:
         #    linear_solution = np.insert(linear_solution, 0, 0, axis=0)
-
+    """
     # number of centers
     Neta = len(eta)
 
@@ -109,7 +113,7 @@ def _NumericalFluxVector(uL, uR):
     # initialise
     lapsed_time = 0
     counter     = 0
-
+    """
     # time integration
     while lapsed_time < tf:
 
@@ -127,60 +131,183 @@ def _NumericalFluxVector(uL, uR):
         F[1:Neta] = _NumericalFluxVector(solution[-1][0:-1], solution[-1][1:])
 
         # compute the left boundary flux
-        if solution[-1][0] < 0.0:
-            uL = 2.0 * solution[-1][0] - solution[-1][1]
-        else:
-            uL = solution[0][0]
-
+        uL   = solution[-1][-1]
         uR   = solution[-1][0]
         F[0] = _NumericalFluxVector(uL, uR)
 
         # compute the right boundary flux
-        if solution[-1][Neta - 1] > 0.0:
-            uR = 2.0 * solution[-1][Neta - 1] - solution[-1][Neta - 2]
-        else:
-            uR = solution[0][Neta - 1]
+        uR      = solution[-1][0]
+        uL      = solution[-1][Neta - 1]
+        F[Neta] = _NumericalFluxVector(uL, uR)
+
+        solution.append(solution[-1][0:Neta] - dt / deta * (F[1:Neta+1] - F[0:Neta]))
+
+        # Eq. (18) Bollati et al. 2021
+        eta_minus = eta_tilde - np.sqrt(2 * C * tf)
+        eta_plus = eta_tilde + np.sqrt(2 * C * tf)
+
+        if lapsed_time >= t_edge:
+            break
+    """
+    # time integration until boundary of the disc
+    #print(t_edge)
+    while lapsed_time <= t_edge:
+
+        # increment
+        counter += 1
+
+        # time step calculation
+        dt = min(deta * CFL / (max(abs(solution[-1])) + 1e-8), 0.02)
 
+        # update time
+        time.append(time[-1] + dt)
+        lapsed_time += dt
+
+        # compute the interior fluxes
+        F[1:Neta] = _NumericalFluxVector(solution[-1][0:-1], solution[-1][1:])
+
+        # compute the left boundary flux
+        uL   = solution[-1][-1]
+        uR   = solution[-1][0]
+        F[0] = _NumericalFluxVector(uL, uR)
+
+        # compute the right boundary flux
+        uR      = solution[-1][0]
         uL      = solution[-1][Neta - 1]
         F[Neta] = _NumericalFluxVector(uL, uR)
 
         solution.append(solution[-1][0:Neta] - dt / deta * (F[1:Neta+1] - F[0:Neta]))
 
+        # Implementing condition for exiting the loop in the asymptotic limit, when N wave and numerical solutions coincide up to 5%
+        # Eq. (18) Bollati et al. 2021.
+        # Left and right extrema of N wave profile
+        eta_minus = eta_tilde - np.sqrt(2 * C * (time[-1] - t0))
+        eta_plus  = eta_tilde + np.sqrt(2 * C * (time[-1] - t0))
+
+        # Maximum and minimum values of N wave Eq. (17) Bollati et al. 2021.
+        N_wave_max = (eta_plus  - eta_tilde) / (time[-1]-t0)
+        N_wave_min = (eta_minus - eta_tilde) / (time[-1]-t0)
+        # Conditions on minimum and maximum amplitude up to 5%
+        condition_max_amplitude = np.abs((N_wave_max - np.max(solution[-1])) / np.max(solution[-1])) <= 0.05
+        condition_min_amplitude = np.abs((N_wave_min - np.min(solution[-1])) / np.min(solution[-1])) <= 0.05 #to be improved, now works only for outer disc. Or not?
+        condition_amplitude     = np.logical_and(condition_max_amplitude,condition_min_amplitude)
+
+        # Conditions on left and right extrema up to 5%
+        condition_max_width = np.abs(eta_plus  - eta[np.argmax(solution[-1])]) <= 0.05
+        condition_min_width = np.abs(eta_minus - eta[np.argmin(solution[-1])]) <= 0.05
+        condition_width     = np.logical_and(condition_max_width, condition_min_width)
+
+
+        if lapsed_time >= tf:
+            if np.logical_and(condition_amplitude, condition_width):
+                break
+
     solution = np.array(solution).transpose()
     time     = np.array(time)
-
-    # plot Fig. 3 Bollati et al. 2021
-    #plt.plot(eta, solution[:,0], label = "$t=t_0+$ ")#+str(0*dt))
-    #plt.plot(eta, solution[:,100], label = "$t=t_0+$ ")#+str(round(dt*Nt/10,2)))
-    #plt.plot(eta, solution[:,200], label = "$t=t_0+$ ")#+str(round(dt*Nt/6,2)))
-    #plt.plot(eta, solution[:,300], label = "$t=t_0+$ ")#+str(round(dt*Nt/3,2)))
-    #plt.plot(eta, solution[:,400], label = "$t=t_0+$ ")#+str(round(T,2)))
-    #plt.legend()
-    #plt.xlabel(r"$\eta$")
-    #plt.ylabel(r"$\chi(t,\eta)$")
-    #plt.grid(True)
-    #plt.title(r'$\chi$ "evolution" $r > r_p$')
-    #plt.show()
-
-    #if show_teta: # combining linear and non-linear solution and plotting
-    #    
-    #    # scale linear solution
-    #    linear_solution = linear_solution * (gamma + 1) * beta_p / 2**(3 / 4)
-    #
-    #    # add linear solution in (t,eta) to non-linear solution array
-    #    total_solution = np.concatenate((linear_solution, solution), axis=1)
-    #    total_time     = np.concatenate((linear_t,        time + t0))
-    #
-    #    fig, ax = plt.subplots(1)
-    #    cont = ax.contourf(total_time, eta, total_solution, levels=np.arange(-4, 4, 0.05), cmap='RdBu')
-    #    for c in cont.collections:
-    #        c.set_rasterized(True)
-    #    plt.colorbar(cont, label='$\chi$')
-    #    ax.set_xlim(0,10)
-    #    ax.set_xlabel('$t$')
-    #    ax.set_ylabel('$\eta$')
+   
+    if False:
+        #plots for debugging: plot \eta profiles to check the evolution of the solution
+        #colorbar
+        n = len(time)
+        norm = cl.Normalize(vmin=time.min(), vmax=time.max())
+        cmap = cm.ScalarMappable(norm=norm, cmap=cm.viridis)
+        cmap.set_array([])
+        
+        plt.figure(figsize=(15,5))
+        
+        idxp = int(n/5) #index to plot 5 profiles equidistant in \t
+        plt.plot(eta, solution[:,0*idxp], label = "$t=t_0+%.1lf$ "%time[0*idxp],
+                color = cmap.to_rgba(time[0*idxp] + 1), ls='-.')#+str(0*dt))
+        plt.plot(eta, solution[:,1*idxp], label = "$t=t_0+%.1lf$ "%time[1*idxp],
+                color = cmap.to_rgba(time[1*idxp] + 1), ls='-.')#+str(round(dt*Nt/10,2)))
+        plt.plot(eta, solution[:,2*idxp], label = "$t=t_0+%.1lf$ "%time[2*idxp],
+                color = cmap.to_rgba(time[2*idxp] + 1), ls='-.')#+str(round(dt*Nt/6,2)))
+        plt.plot(eta, solution[:,3*idxp], label = "$t=t_0+%.1lf$ "%time[3*idxp],
+                color = cmap.to_rgba(time[3*idxp] + 1), ls='-.')#+str(round(dt*Nt/3,2)))
+        plt.plot(eta, solution[:,4*idxp], label = "$t=t_0+%.1lf$ "%time[4*idxp],
+                color = cmap.to_rgba(time[4*idxp] + 1), ls='-.')#+str(round(T,2)))
+        plt.plot(eta, solution[:,-1], label = "$t=t_0+%.1lf$ "%time[-1],
+                color = cmap.to_rgba(time[-1]), ls='-.')#+str(round(T,2)))
+
+        #check last profile with analytic N wave to see if asymptotic limit already reached
+        if eta_tilde <= 0:
+            eta_plus = eta_tilde + np.sqrt(2*C*(time[-1]-t0))
+            eta_minus = eta_tilde - np.sqrt(2*C*(time[-1]-t0))
+            #plt.plot(eta, np.where(np.logical_and(eta>eta_minus,eta<eta_plus),(eta-eta_tilde)/(time[-1]-t0),0), label='N')
+        else:
+            eta_plus = eta_tilde + np.sqrt(2*C*(time[-1]-t0))
+            eta_minus = eta_tilde - np.sqrt(2*C*(time[-1]-t0))
+            #plt.plot(eta, np.where(np.logical_and(eta>eta_minus,eta<eta_plus),-(eta-eta_tilde)/(time[-1]-t0),0), label='N')
+        
+        plt.legend()
+        plt.xlim(-40,40)
+        plt.xlabel(r"$\eta$")
+        plt.ylabel(r"$\chi(t,\eta)$")
+        plt.grid(True)
+        cb = plt.colorbar(cmap, label = r't', shrink = 1)
+        #cb.ax.tick_params(labelsize=20)
+        cb.ax.set_ylabel(r't')#, fontsize=20)
+        cb.ax.yaxis.set_label_position('left')
+        cb.ax.set_aspect('auto')
+        #Trying to get nice double colorbars, failing miserably
+        """
+        cb2 = cb.ax.twinx() 
+        cb2.set_ylim([time[0], time[-1]])
+        
+        rticks=np.zeros(np.shape(time))
+        for i in range(len(time)):
+            if i%idxp == 0 or i == len(time)-1:
+                def invert_t(x, Rp, hr, q, p, m_p, m_th):
+                    return _t_vector(x, Rp, hr, q, p, m_p, m_th) - time[i]
+                rticks[i] = fsolve(invert_t, args=(Rp, hr, q, p, m_p, m_th), x0 = 120)
+        cb2.set_yticks([time[0*idxp],time[1*idxp],time[2*idxp],time[3*idxp],time[4*idxp],time[-1]])
+        cb2.set_yticklabels([round(rticks[0*idxp]),round(rticks[1*idxp]),round(rticks[2*idxp]),
+                            round(rticks[3*idxp]),round(rticks[4*idxp]),round(rticks[-1])])
+        #cb2.set_yticks([tic for tic in pdfp])
+        cb2.set_ylabel(r'R/R$_{\rm p}$', labelpad=8)#, fontsize=20)
+        #cb2.tick_params(labelsize=20)
+        """
+        if eta_tilde <= 0:
+            plt.title(r'$\chi$ "evolution" $r > r_p$')
+        else:
+            plt.title(r'$\chi$ "evolution" $r < r_p$')
+        plt.show()
+        
+        #Like before, check last profile with analytic N wave to see if asymptotic limit already reached
+        plt.plot(eta, solution[:,-1], label = "$t=t_0+%.1lf$ "%time[-1])#+str(round(T,2)))
+        #plt.plot(eta, solution[:,-1], label = "$t=t_0+$ ")#+str(round(T,2)))
+        if eta_tilde <= 0:
+            plt.plot(eta, np.where(np.logical_and(eta>eta_minus,eta<eta_plus),(eta-eta_tilde)/(time[-1]-t0),0), label='N')
+        else:
+            plt.plot(eta, np.where(np.logical_and(eta>eta_minus,eta<eta_plus),(eta-eta_tilde)/(time[-1]-t0),0), label='N')
+        plt.legend()
+        plt.xlabel(r"$\eta$")
+        plt.ylabel(r"$\chi(t,\eta)$")
+        plt.grid(True)
+        #plt.xticks([-28, -27, -26, -25, -24, 0, 20, 21, 22])
+        plt.show()
+        
+    #plt.plot(time, rticks)
+
+    if show_teta: # combining linear and non-linear solution and plotting
+        
+        # scale linear solution
+        linear_solution = linear_solution * (gamma + 1) * beta_p / 2**(3 / 4)
+    
+        # add linear solution in (t,eta) to non-linear solution array
+        total_solution = np.concatenate((linear_solution, solution), axis=1)
+        total_time     = np.concatenate((linear_t,        time + t0))
+    
+        fig, ax = plt.subplots(1)
+        cont = ax.contourf(total_time, eta, total_solution, levels=np.arange(-4, 4, 0.05), cmap='RdBu')
+        for c in cont.collections:
+            c.set_rasterized(True)
+        plt.colorbar(cont, label=r'$\chi$')
+        ax.set_xlim(0,10)
+        ax.set_xlabel(r'$t$')
+        ax.set_ylabel(r'$\eta$')
     #    #plt.savefig("teta_badjoin.pdf")
-    #    plt.show()
+        plt.show()
 
     # The following will put the linear solution into returned solution for Chi
     #if False:
@@ -195,4 +322,4 @@ def _NumericalFluxVector(uL, uR):
     #    return total_time, eta, total_solution, eta_inner, solution_inner
 
     return time, eta, solution
-    
\ No newline at end of file
+    
diff --git a/src/wakeflow/data/Basement/global_linear_perturbations_wrong.npy.tar.gz b/src/wakeflow/data/Basement/global_linear_perturbations_wrong.npy.tar.gz
new file mode 100644
index 0000000..c30c6ab
Binary files /dev/null and b/src/wakeflow/data/Basement/global_linear_perturbations_wrong.npy.tar.gz differ
diff --git a/src/wakeflow/data/global_linear_perturbations.npy.tar.gz b/src/wakeflow/data/global_linear_perturbations.npy.tar.gz
new file mode 100644
index 0000000..d98bd5a
Binary files /dev/null and b/src/wakeflow/data/global_linear_perturbations.npy.tar.gz differ
diff --git a/src/wakeflow/data/global_linear_perturbations_mesh.npy.tar.gz b/src/wakeflow/data/global_linear_perturbations_mesh.npy.tar.gz
new file mode 100644
index 0000000..de7dd8d
Binary files /dev/null and b/src/wakeflow/data/global_linear_perturbations_mesh.npy.tar.gz differ
diff --git a/src/wakeflow/discminer_interface.py b/src/wakeflow/discminer_interface.py
new file mode 100644
index 0000000..3feda76
--- /dev/null
+++ b/src/wakeflow/discminer_interface.py
@@ -0,0 +1,81 @@
+# discminer_interface.py
+# Written by Daniele Fasano, based on wakefit_interface.py by Thomas Hilder
+
+"""
+Contains the Discminer class, intended for use in Discminer to generate wake models
+"""
+
+import subprocess, os
+from .model_setup         import _run_setup, _Parameters
+from .grid                import _Grid
+from .linear_perts        import _LinearPerts
+from .non_linear_perts    import _NonLinearPerts
+from .wakeflow            import WakeflowModel
+
+# class for wakefit to use to generate velocity perturbations for models
+class _WakeminerModel(WakeflowModel):
+
+    # generate the model using the configuration specified by the user
+    def run(self) -> None:
+
+        # check requirements for models for wakefit are satisfied (don't make files!)
+        assert self.model_params["make_midplane_plots"] == False
+        assert self.model_params["save_total"]          == False
+        assert self.model_params["save_perturbations"]  == False
+        assert type(self.model_params["m_planet"])      is not list
+        # run setup
+        try:
+            params = _run_setup(self.model_params)
+        except ArithmeticError:
+            raise Exception("Model has not been configured.")
+
+        # run wakeflow using custom method for discminer
+        return self._run_wakeflow_for_discminer(params)
+
+    # internal method that is called by self.run to generate the results for a specific set of parameters
+    def _run_wakeflow_for_discminer(self, params: _Parameters) -> None:
+        """
+        Internal use method for generating the planet wake by calling other parts of Wakeflow.
+
+        Parameters
+        ----------
+        params : Parameters 
+            Paramaters object specfiying options for the model to be generated.
+        """
+
+        print("Extracting linear perturbations nearby planet")
+
+        # make empty grid for linear perturbations
+        grid_lin_perts = _Grid(params)
+        grid_lin_perts._make_grid()
+        grid_lin_perts._make_empty_disk()
+
+        # extract linear perturbations from file
+        lin_perts = _LinearPerts(params)
+        lin_perts._cut_annulus_segment()
+
+        # add the linear perturbations onto grid
+        grid_lin_perts._add_linear_perturbations(lin_perts, rho_background=0.0)
+
+        # make empty grid for non-linear perturbations
+        grid_nonlin_perts = _Grid(params)
+        grid_nonlin_perts._make_grid()
+        grid_nonlin_perts._make_empty_disk()
+
+        # initialise non-linear perturbations
+        nonlin_perts = _NonLinearPerts(params, grid_nonlin_perts)
+
+        # extract initial condition from the linear perturbations
+        nonlin_perts._extract_ICs_ann(lin_perts)
+
+        # solve for non-linear perturbations
+        nonlin_perts._get_non_linear_perts()
+
+        # add non-linear perturbations to grid
+        grid_nonlin_perts._add_non_linear_perturbations(nonlin_perts, rho_background=0.0)
+
+        # merge grids for result
+        grid_lin_perts._merge_grids(grid_nonlin_perts)
+
+        # return the grid and velocity components
+        return (grid_lin_perts.x, grid_lin_perts.y, grid_lin_perts.v_r, grid_lin_perts.v_phi)
\ No newline at end of file
diff --git a/src/wakeflow/grid.py b/src/wakeflow/grid.py
index f64560a..95b5f96 100644
--- a/src/wakeflow/grid.py
+++ b/src/wakeflow/grid.py
@@ -122,9 +122,20 @@ def _make_cartesian_grid(self) -> None:
         """
         
         # make grid from specified parameters
-        self.x    = np.linspace(-self.p.r_outer, self.p.r_outer, self.p.n_x)
-        self.y    = np.linspace(-self.p.r_outer, self.p.r_outer, self.p.n_y)
-        self.z_xy = np.linspace(0, self.height, self.p.n_z)
+        self.x, stepx_    = np.linspace(-self.p.r_outer, self.p.r_outer, self.p.n_x, retstep="True")
+        self.y, stepy_    = np.linspace(-self.p.r_outer, self.p.r_outer, self.p.n_y, retstep="True")
+        self.z_xy         = np.linspace(0, self.height, self.p.n_z)
+
+        # extending the grid by a sqrt(2) factor to avoid interpolations effects at the corner when rotating the grid to get the planet position, using the same spacing of the original grid
+        if self.p.rot_interp == True:
+            x_grid_ext_l = np.arange(-np.sqrt(2)*self.p.r_outer, -self.p.r_outer, stepx_)
+            y_grid_ext_l = np.arange(-np.sqrt(2)*self.p.r_outer, -self.p.r_outer, stepy_)
+            
+            x_grid_ext_r = -x_grid_ext_l[::-1]
+            y_grid_ext_r = -y_grid_ext_l[::-1]
+            
+            self.x = np.concatenate([x_grid_ext_l, self.x, x_grid_ext_r])
+            self.y = np.concatenate([y_grid_ext_l, self.y, y_grid_ext_r])
 
         self.X, self.Z_xy, self.Y  = np.meshgrid(self.x, self.z_xy, self.y, indexing='ij')
 
@@ -345,6 +356,29 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
             unperturbed density from Grid object where make_keplerian_disk has been used.
         """
 
+        # === Daniele code === #
+
+        # get linear solution           
+        lp = LinearPerts
+
+        # segment radial size (in units of Hill radius), note for conversions that self.p.l = 1 Hill radius in cgs.
+        r_box_size_left = lp.x_box_left
+        r_box_size_right = lp.x_box_right
+        # segment azimuthal size (in units of \pi). 
+        phi_box_size_top = lp.x_box_top / 2
+        phi_box_size_bottom = lp.x_box_bottom / 2
+        # segment radial extrema for masking
+        min_r = self.p.r_planet - r_box_size_left * self.p.l
+        max_r = self.p.r_planet + r_box_size_right * self.p.l
+        # segment azimuthal extrema for masking
+        max_phi =  phi_box_size_top * np.pi 
+        min_phi = -phi_box_size_bottom * np.pi 
+        
+        # === #
+
+        """
+        # === Tom code === #
+        
         # box size (in units of Hill radius), note for conversions that self.p.l = 1 Hill radius in cgs
         x_box_size_l = 2 * self.p.scale_box_l
         x_box_size_r = 2 * self.p.scale_box_r
@@ -369,6 +403,9 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
         else:
             max_phi = np.arctan2(max_y, max_r) #+ (np.pi / 2)
             min_phi = np.arctan2(min_y, max_r) #+ (np.pi / 2)
+            
+        # === #
+        """
 
         # find (phi, r) grid for either Cartesian or Cylindrical global grid
         if self.info["Type"] == "cartesian":
@@ -385,6 +422,17 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
         
         # new PHI grid to use (-pi,pi) instead of (0,2pi), where values are swapped in place, also ditch z coordinate
         # also construct a mask that contains 0 outside linear annulus and 1 inside
+        # === Daniele code === #
+  
+        R_new       = R[:,0,:]
+        PHI_new     = np.where(PHI[:,0,:]>np.pi, PHI[:,0,:] - 2*np.pi, PHI[:,0,:])
+        linear_mask = np.where(np.logical_and(np.logical_and(PHI_new>=min_phi,PHI_new<=max_phi), np.logical_and(R_new>min_r,R_new<max_r)), 1, 0)
+        
+        # === #
+
+        """
+        # === Tom code === #
+
         PHI_new     = np.zeros((PHI.shape[0],PHI.shape[2]))
         PHI_new_p   = np.zeros((PHI.shape[0],PHI.shape[2]))
         linear_mask = np.zeros((PHI.shape[0],PHI.shape[2]))
@@ -424,6 +472,9 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
                         linear_mask[i,j] = 1
                     else:
                         linear_mask[i,j] = 0
+        
+        # === #
+        """
 
         # get linear solution           
         lp = LinearPerts
@@ -443,6 +494,17 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
                 PHI_flat_new[i] = PHI_flat[i]
         
         # assemble interpolation functions over linear perts grid
+        # === Daniele code === #
+
+        interp_v_r   = RectBivariateSpline(lp.phi_ann, lp.r_ann, lp.pert_v_r_ann)
+        interp_v_phi = RectBivariateSpline(lp.phi_ann, lp.r_ann, lp.pert_v_phi_ann)
+        interp_rho   = RectBivariateSpline(lp.phi_ann, lp.r_ann, lp.pert_rho_ann)
+        
+        # === #
+
+        """
+        # === Tom code === #
+        
         interp_v_r   = LinearNDInterpolator(np.transpose([PHI_flat_new, R_flat]), lp.pert_v_r_ann  .flatten())
         interp_v_phi = LinearNDInterpolator(np.transpose([PHI_flat_new, R_flat]), lp.pert_v_phi_ann.flatten())
         interp_rho   = LinearNDInterpolator(np.transpose([PHI_flat_new, R_flat]), lp.pert_rho_ann  .flatten())
@@ -454,6 +516,10 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
         PHI_planet_new = np.zeros((PHI_planet.shape[0],PHI_planet.shape[1]))
         for i in range(PHI_planet.shape[0]):
             for j in range(PHI_planet.shape[1]):
+        
+        # === #
+        """
+
 
                 # transforming phi coordinate in place
                 if PHI_planet[i,j] > np.pi:
@@ -462,6 +528,18 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
                     PHI_planet_new[i,j] = PHI_planet[i,j]
         
         # evaluate interpolations on global grid
+        
+        # === Daniele code === #
+
+        global_lin_v_r   = interp_v_r.ev  (PHI_new, R_new)
+        global_lin_v_phi = interp_v_phi.ev(PHI_new, R_new)
+        global_lin_rho   = interp_rho.ev  (PHI_new, R_new)
+        
+        # === #
+
+        """
+        # === Tom code === #
+
         global_lin_v_r   = np.nan_to_num(interp_v_r  (PHI_planet_new, R_new))
         global_lin_v_phi = np.nan_to_num(interp_v_phi(PHI_planet_new, R_new))
         global_lin_rho   = np.nan_to_num(interp_rho  (PHI_planet_new, R_new))
@@ -471,6 +549,9 @@ def _add_linear_perturbations(self, LinearPerts : _LinearPerts, rho_background :
         #plt.title("first one")
         #plt.imshow(global_lin_v_r, cmap="RdBu", origin="lower")
         #plt.show()
+        
+        # === #
+        """
 
         # apply mask to only get solution in valid domain. This transpose business is confusing but works
         if self.p.grid_type == "cartesian":
@@ -648,7 +729,7 @@ def _smooth_box(self, big_box_grid : "_Grid") -> None:
 #            return False
 #
 #        # merge data arrays
-#        self.rho = g.rho
+#        self.rho = g.rho 
 
     # create plots of a constant z slice, mostly used for debugging but also shows midplane results to user
     # if they set show_midplane_plots=True
@@ -699,6 +780,13 @@ def _show_disk2D(
             plt.close("all")
             fig, ax = plt.subplots(dpi=150)
             c       = ax.pcolormesh(self.x, self.y, np.transpose(self.v_r[:,z_slice,:]), vmin=-vr_max, vmax=vr_max, cmap='RdBu', rasterized=True)
+            if False:
+                from .transformations import _phi_wake
+                r_wake = np.linspace(5, self.p.r_outer, 1000)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, True)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='r', linewidth=0.5)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, False)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='k', linewidth=0.5)
             ax.axis('scaled')
             ax.set_title(r"$\delta v_r$")
             if not dimless:
@@ -718,6 +806,13 @@ def _show_disk2D(
             plt.close("all")
             fig, ax = plt.subplots(dpi=150)
             c       = ax.pcolormesh(self.x, self.y, np.transpose(self.v_phi[:,z_slice,:]), vmin=-vphi_max, vmax=vphi_max, cmap='RdBu', rasterized=True)
+            if False:
+                from .transformations import _phi_wake
+                r_wake = np.linspace(5, self.p.r_outer, 1000)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, True)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='r', linewidth=0.5)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, False)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='k', linewidth=0.5)
             ax.axis('scaled')
             ax.set_title(r"$\delta v_{\phi}$")
             if not dimless:
@@ -737,6 +832,13 @@ def _show_disk2D(
             plt.close("all")
             fig, ax = plt.subplots(dpi=150)
             c       = ax.pcolormesh(self.x, self.y, np.transpose(self.rho[:,z_slice,:]), vmin=-rho_max, vmax=rho_max, cmap='RdBu', rasterized=True)
+            if False:
+                from .transformations import _phi_wake
+                r_wake = np.linspace(5, self.p.r_outer, 1000)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, True)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='r', linewidth=0.5)
+                wake_ = _phi_wake(r_wake, self.p.r_planet, self.p.hr, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, False)
+                ax.plot(r_wake*np.cos(wake_), r_wake*np.sin(wake_), ls = '--', c='k', linewidth=0.5)
             ax.axis('scaled')
             ax.set_title(r"$\delta \rho \, / \rho$")
             if not dimless:
@@ -907,4 +1009,5 @@ def _save_results(self, label : str, printed : str) -> None:
 
         print(f"{printed} saved to {savedir}")
 
-        
+
+
diff --git a/src/wakeflow/linear_perts.py b/src/wakeflow/linear_perts.py
index ca51247..ea627c8 100644
--- a/src/wakeflow/linear_perts.py
+++ b/src/wakeflow/linear_perts.py
@@ -12,6 +12,7 @@
 import numpy                    as np
 import matplotlib.pyplot        as plt
 from scipy.interpolate      import RectBivariateSpline
+from .transformations       import _Eta_vector
 
 if TYPE_CHECKING:
     from .model_setup       import _Parameters
@@ -29,67 +30,208 @@ def __init__(self, parameters : '_Parameters') -> None:
 
         # grab parameters object
         self.p = parameters
-
-        # get location of linear perturbations data files
-        pert_loc = pkg_resources.resource_filename('wakeflow', 'data/linear_perturbations.npy')
-        mesh_loc = pkg_resources.resource_filename('wakeflow', 'data/linear_perturbations_mesh.npy')
-
-        # read perturbations from files
-        try:
-            perts = np.load(pert_loc)
-            mesh  = np.load(mesh_loc)
-
-        # in the case files have not been extracted from tarballs yet, extract them
-        except FileNotFoundError:
-
-            # open tarballs
-            pert_tar = tarfile.open(f"{pert_loc}.tar.gz")
-            mesh_tar = tarfile.open(f"{mesh_loc}.tar.gz")
-
-            # extract npy files
-            loc = pkg_resources.resource_filename('wakeflow', 'data')
-            pert_tar.extractall(loc)
-            mesh_tar.extractall(loc)
-
-            # close tarballs
-            pert_tar.close()
-            mesh_tar.close()
-
-            # read perturbations from files
-            perts = np.load(pert_loc)
-            mesh  = np.load(mesh_loc)
-
-        # get perturbation arrays
-        self.pert_v_r   = perts[0]
-        self.pert_v_phi = perts[1]
-        self.pert_rho   = perts[2]
-
-        # grid
-        self.X = mesh[0]
-        self.Y = mesh[1]
-
-        # linear perturbations read in grid
-        x = self.X[0,:]
-        y = self.Y[:,0]
-
-        # define constants for linear perts
-        self.x_box_l = 2 * self.p.scale_box_l
-        self.x_box_r = 2 * self.p.scale_box_r
-        self.x_box_t = 2 * self.p.scale_box_ang_t
-        self.x_box_b = 2 * self.p.scale_box_ang_b
         
-        # cut square box grid in linear regime
-        self.x_cut = x[np.argmin(x < -self.x_box_l) : np.argmin(x < self.x_box_r) + 1]
-        self.y_cut = y[np.argmin(y < -self.x_box_b) : np.argmin(y < self.x_box_t) + 1]
-
-        # test plot
-        if False:
-            plt.figure()
-            plt.contourf(x, y, self.pert_rho, levels=100, cmap='RdBu')
-            plt.xlim(self.x_cut[0],self.x_cut[-1])
-            plt.ylim(self.y_cut[0],self.y_cut[-1])
-            plt.colorbar()
-            plt.show()
+        # initialise perturbations properties
+        self.info = {
+            "Type": None,
+            "Grid": None,
+            "Size": [0, 0, 0]
+            }
+        
+        # Using linear perturbations computed following Bollati et al. 2021 with shearing sheet assumption
+        # NOTE: this is deprecated and will probably lead to wrong results
+        if self.p.lin_type == "shearing_sheet":
+            # get location of linear perturbations data files
+            pert_loc = pkg_resources.resource_filename('wakeflow', 'data/linear_perturbations.npy')
+            mesh_loc = pkg_resources.resource_filename('wakeflow', 'data/linear_perturbations_mesh.npy')
+    
+            # read perturbations from files
+            try:
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+    
+            # in the case files have not been extracted from tarballs yet, extract them
+            except FileNotFoundError:
+    
+                # open tarballs
+                pert_tar = tarfile.open(f"{pert_loc}.tar.gz")
+                mesh_tar = tarfile.open(f"{mesh_loc}.tar.gz")
+    
+                # extract npy files
+                loc = pkg_resources.resource_filename('wakeflow', 'data')
+                pert_tar.extractall(loc)
+                mesh_tar.extractall(loc)
+    
+                # close tarballs
+                pert_tar.close()
+                mesh_tar.close()
+    
+                # read perturbations from files
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+    
+            # get perturbation arrays
+            self.pert_v_r   = perts[0]
+            self.pert_v_phi = perts[1]
+            self.pert_rho   = perts[2]
+    
+            # grid
+            self.X = mesh[0]
+            self.Y = mesh[1]
+    
+            # linear perturbations read in grid
+            x = self.X[0,:]
+            y = self.Y[:,0]
+    
+            # define constants for linear perts
+            self.x_box_left   = 2 * self.p.scale_box_left
+            self.x_box_right  = 2 * self.p.scale_box_right
+            self.x_box_top    = 2 * self.p.scale_box_ang_top
+            self.x_box_bottom = 2 * self.p.scale_box_ang_bottom
+            
+            # cut square box grid in linear regime
+            self.x_cut = x[np.argmin(x < -self.x_box_l) : np.argmin(x < self.x_box_r) + 1]
+            self.y_cut = y[np.argmin(y < -self.x_box_b) : np.argmin(y < self.x_box_t) + 1]
+    
+            # test plot
+            if False:
+                plt.figure()
+                plt.contourf(x, y, self.pert_rho, levels=100, cmap='RdBu')
+                plt.xlim(self.x_cut[0],self.x_cut[-1])
+                plt.ylim(self.y_cut[0],self.y_cut[-1])
+                plt.colorbar()
+                plt.show()
+
+            #updating info of linear perts        
+            self.info["Type"]    = "shearing_sheet"
+            self.info["Grid"]    = "cartesian"
+            self.info["Size"][0] = x.shape[0]
+            self.info["Size"][1] = y.shape[0]
+            self.info["Size"][2] = self.p.n_z
+
+        #Using global linear perturbations computed using Miranda's code.         
+        elif self.p.lin_type == "global":
+            # get location of linear perturbations data files
+            pert_loc = pkg_resources.resource_filename('wakeflow', 'data/global_linear_perturbations.npy')
+            mesh_loc = pkg_resources.resource_filename('wakeflow', 'data/global_linear_perturbations_mesh.npy')
+    
+            # read perturbations from files
+            try:
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+    
+            # in the case files have not been extracted from tarballs yet, extract them
+            except FileNotFoundError:
+    
+                # open tarballs
+                pert_tar = tarfile.open(f"{pert_loc}.tar.gz")
+                mesh_tar = tarfile.open(f"{mesh_loc}.tar.gz")
+    
+                # extract npy files
+                loc = pkg_resources.resource_filename('wakeflow', 'data')
+                pert_tar.extractall(loc)
+                mesh_tar.extractall(loc)
+    
+                # close tarballs
+                pert_tar.close()
+                mesh_tar.close()
+    
+                # read perturbations from files
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+                
+            # get perturbation arrays
+            self.pert_v_r   = perts[0]
+            self.pert_v_phi = perts[1]
+            self.pert_rho   = perts[2]
+    
+            # grid
+            self.R   = mesh[0]
+            self.PHI = mesh[1]
+
+            #Rescale R with planet radius
+            self.R *= self.p.r_planet
+    
+            # linear perturbations read in grid
+            r   = self.R[0,:]
+            phi = self.PHI[:,0]
+
+            #defining cartesian grid
+            x = np.linspace(-np.max(r), np.max(r), len(r))
+            y = np.linspace(-np.max(r), np.max(r), len(r))
+
+            #creating cartesian mesh
+            self.X, self.Y = np.meshgrid(x, y)
+
+            # define constants for linear perts
+            self.x_box_left   = 2 * self.p.scale_box_left
+            self.x_box_right  = 2 * self.p.scale_box_right
+            self.x_box_top    = 2 * self.p.scale_box_ang_top
+            self.x_box_bottom = 2 * self.p.scale_box_ang_bottom
+
+            #updating info of linear perts      
+            self.info["Type"]    = "global"
+            self.info["Grid"]    = "cylindrical"
+            self.info["Size"][0] = r.shape[0]
+            self.info["Size"][1] = phi.shape[0]
+            self.info["Size"][2] = self.p.n_z
+
+        #Using perturbations from hydro simulation (not yet implemented)
+        elif self.p.lin_type == "simulation":
+            # get location of linear perturbations data files
+            pert_loc = pkg_resources.resource_filename('wakeflow', 'data/simulation_linear_perturbations.npy')
+            mesh_loc = pkg_resources.resource_filename('wakeflow', 'data/simulation_linear_perturbations_mesh.npy')
+    
+            # read perturbations from files
+            try:
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+    
+            # in the case files have not been extracted from tarballs yet, extract them
+            except FileNotFoundError:
+    
+                # open tarballs
+                pert_tar = tarfile.open(f"{pert_loc}.tar.gz")
+                mesh_tar = tarfile.open(f"{mesh_loc}.tar.gz")
+    
+                # extract npy files
+                loc = pkg_resources.resource_filename('wakeflow', 'data')
+                pert_tar.extractall(loc)
+                mesh_tar.extractall(loc)
+    
+                # close tarballs
+                pert_tar.close()
+                mesh_tar.close()
+    
+                # read perturbations from files
+                perts = np.load(pert_loc)
+                mesh  = np.load(mesh_loc)
+                
+            # get perturbation arrays
+            self.pert_v_r   = perts[0]
+            self.pert_v_phi = perts[1]
+            self.pert_rho   = perts[2]
+    
+            # grid
+            self.X = mesh[0]
+            self.Y = mesh[1]
+    
+            # linear perturbations read in grid
+            x = self.X[0,:]
+            y = self.Y[:,0]
+
+            # define constants for linear perts
+            self.x_box_left   = 2 * self.p.scale_box_left
+            self.x_box_right  = 2 * self.p.scale_box_right
+            self.x_box_top    = 2 * self.p.scale_box_ang_top
+            self.x_box_bottom = 2 * self.p.scale_box_ang_bottom
+
+            #updating info of linear perts      
+            self.info["Type"]    = "simulation"
+            self.info["Grid"]    = "cartesian"
+            self.info["Size"][0] = self.x.shape[0]
+            self.info["Size"][1] = self.y.shape[0]
+            self.info["Size"][2] = self.p.n_z
 
     # old method of extracting linear perturbations 
     def _cut_box_square(self) -> None:
@@ -97,7 +239,7 @@ def _cut_box_square(self) -> None:
         """
 
         # box size (in units of Hill radius), with default scale_box = 1. (note for conversions that self.p.l = 1 Hill radius in cgs)
-        box_size = 2*self.p.scale_box_l
+        box_size = 2*self.p.scale_box_left
         artificial_y_scale = 6
 
         # linear perturbations read in grid
@@ -141,127 +283,201 @@ def _cut_box_square(self) -> None:
         self.y_cut = y_cut
         self.x_sq  = self.p.l * x_cut + self.p.r_planet
         self.y_sq  = self.p.l * y_cut
-        
+
     # extract linear perturbations and interpolate onto annulus segment grid for global results
-    def _cut_box_annulus_segment(self) -> None:
+    def _cut_annulus_segment(self) -> None:
         """Extract the part of the linear solution needed for the model, and interpolate onto appropriate grid
         """
 
-        # box size (in units of Hill radius), note for conversions that self.p.l = 1 Hill radius in cgs
-        x_box_size_l = 2*self.p.scale_box_l
-        x_box_size_r = 2*self.p.scale_box_r
-        y_box_size_t = 2*self.p.scale_box_ang_t
-        y_box_size_b = 2*self.p.scale_box_ang_b
+        # segment radial size (in units of Hill radius), note for conversions that self.p.l = 1 Hill radius in cgs. 
+        r_box_size_left  = self.x_box_left
+        r_box_size_right = self.x_box_right
+        # segment azimuthal size (in units of \pi). 
+        phi_box_size_top    = self.x_box_top / 2
+        phi_box_size_bottom = self.x_box_bottom / 2
+
+        #interpolate perturbations on a cylindrical grid and evaluate them in the annulus segment
+        if self.info["Type"] == "shearing_sheet" or self.info["Type"] == "simulation":
+
+            # linear perturbations read in grid
+            x = self.X[0,:]
+            y = self.Y[:,0]
+
+            # cut square box in linear regime
+            x_cut = x[np.argmin(x < -r_box_size_left) : np.argmin(x < r_box_size_right) + 1]
+
+            # annulus segment grid, radial granularity from square box, angular granularity fixed for now
+            r = np.linspace(
+                self.p.r_planet - r_box_size_left*self.p.l, 
+                self.p.r_planet + r_box_size_right*self.p.l, 
+                len(x_cut)
+            )
+            
+            phi = np.linspace(
+                -phi_box_size_bottom*np.pi,
+                phi_box_size_top*np.pi,
+                1000
+            )
+
+            R, PHI = np.meshgrid(r, phi)
+            
+            # preparing pertubations for interpolation
+            v_r_cart   = self.pert_v_r   
+            v_phi_cart = self.pert_v_phi 
+            rho_cart   = self.pert_rho
+            
+            #flip perts  if rotation is clockwise
+            if self.p.a_cw == -1:
+                v_r_cart   =  np.flipud(v_r_cart)
+                v_phi_cart = -np.flipud(v_phi_cart)
+                rho_cart   =  np.flipud(rho_cart)
+                
+            # interpolation over global linear (cartesian) grid
+            interp_v_r   = RectBivariateSpline(y, x, v_r_cart)
+            interp_v_phi = RectBivariateSpline(y, x, v_phi_cart)
+            interp_v_rho = RectBivariateSpline(y, x, rho_cart)
+
+            #evaluation of cartesian coordinates corresponding to our polar grid
+            X_pert_grid = (R * np.cos(PHI) - self.p.r_planet)/self.p.l
+            Y_pert_grid = R * np.sin(PHI)/self.p.l
+
+            # evaluate interpolation over our annulus segment
+            self.pert_v_r_ann   = interp_v_r.ev  (Y_pert_grid, X_pert_grid)
+            self.pert_v_phi_ann = interp_v_phi.ev(Y_pert_grid, X_pert_grid)
+            self.pert_rho_ann   = interp_v_rho.ev(Y_pert_grid, X_pert_grid)
+
+            #plot for debugging
+            if False:
+                plt.imshow(self.pert_rho_ann, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r[0], r[-1], phi[0], phi[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+            if False:
+                eta = _Eta_vector(self.p.r_planet + r_box_size_left*self.p.l, phi, self.p.r_planet, self.p.hr, self.p.q, self.p.p, self.p.cw_rotation, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+                plt.plot(phi, self.pert_rho_ann[:,-1])
+                ax = plt.gca()
+                ax.set_xlabel(r'$\varphi$ [rad]')
+                ax.set_ylabel(r'$\sigma$')
+                ax2 = ax.twiny()
+                ax2.plot(eta, self.pert_rho_ann[:,-1])
+                #ax2.plot(eta, self.pert_rho_ann[:,np.argmin(r_<self.p.r_planet + x_box_size_r*self.p.l)])
+                ax2.set_xlabel(r'$\eta$')
+                plt.show()
+
+            # scale to cgs units
+            self.pert_v_r_ann   *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
+            self.pert_v_phi_ann *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
+            self.pert_rho_ann   *=                   (self.p.m_planet/self.p.m_thermal)
+
+            # scale velocities to km/s
+            self.pert_v_r_ann   *= 1e-5
+            self.pert_v_phi_ann *= 1e-5
+
+            # save annulus grid
+            self.r_ann   = r
+            self.phi_ann = phi
+            self.R_ann   = R
+            self.PHI_ann = PHI
+
+        #The global solution is computed on a cylindrical grid, no need to interpolate
+        elif self.info["Type"] == "global":
+
+            # linear perturbations read in grid
+            r   = self.R[0,:]
+            phi = self.PHI[:,0]
+
+            #masks for restriction
+            r_mask = np.logical_and(r >= self.p.r_planet - r_box_size_left*self.p.l, r <= self.p.r_planet + r_box_size_right*self.p.l)
+            phi_mask = np.logical_and(phi >= -phi_box_size_bottom*np.pi, phi <= phi_box_size_top*np.pi)
+            #print(r_mask.shape, phi_mask.shape)
+            #restricting grid to the annulus segment
+            r_ann   = r[r_mask]
+            phi_ann = phi[phi_mask]
+
+            R_ann, PHI_ann = np.meshgrid(r_ann, phi_ann)
+
+            # preparing pertubations for annulus segment
+            v_r_cyl   = self.pert_v_r.T   
+            v_phi_cyl = self.pert_v_phi.T
+            rho_cyl   = self.pert_rho.T
+
+            #plot for debugging
+            if False:
+                plt.imshow(rho_cyl, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+                plt.imshow(v_r_cyl, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+                plt.imshow(v_phi_cyl, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+            #print(v_r_cyl.shape)
+
+            #flip perts  if rotation is clockwise
+            if self.p.a_cw == -1:
+                v_r_cyl   =  np.flipud(v_r_cyl)
+                v_phi_cyl = -np.flipud(v_phi_cyl)
+                rho_cyl   =  np.flipud(rho_cyl)
+
+            #restricting perturbations to the annulus segment
+            self.pert_v_r_ann   = v_r_cyl[:,r_mask][phi_mask,:]
+            self.pert_v_phi_ann = v_phi_cyl[:,r_mask][phi_mask,:]
+            self.pert_rho_ann   = rho_cyl[:,r_mask][phi_mask,:]
+
+            #plot for debugging
+            if False:
+                plt.imshow(self.pert_rho_ann, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+                plt.imshow(self.pert_v_r_ann, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+                plt.imshow(self.pert_v_phi_ann, cmap="RdBu", vmin=-1, vmax=1, origin='lower', extent=(r_ann[0], r_ann[-1], phi_ann[0], phi_ann[-1]))
+                plt.axis('auto')
+                plt.xlabel('R [au]')
+                plt.ylabel(r'$\varphi$ [rad]')
+                plt.show()
+            if False:
+                eta = _Eta_vector(self.p.r_planet + r_box_size_left*self.p.l, phi_ann, self.p.r_planet, self.p.hr, self.p.q, self.p.p, self.p.cw_rotation, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+                plt.plot(phi_ann, self.pert_rho_ann[:,-1])
+                ax = plt.gca()
+                ax.set_xlabel(r'$\varphi$ [rad]')
+                ax.set_ylabel(r'$\sigma$')
+                ax2 = ax.twiny()
+                ax2.plot(eta, self.pert_rho_ann[:,-1])
+                #ax2.plot(eta, self.pert_rho_ann[:,np.argmin(r_<self.p.r_planet + x_box_size_r*self.p.l)])
+                ax2.set_xlabel(r'$\eta$')
+                plt.show()
+
+            # scale to cgs units
+            self.pert_v_r_ann   *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
+            self.pert_v_phi_ann *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
+            self.pert_rho_ann   *=                   (self.p.m_planet/self.p.m_thermal)
+
+            # scale velocities to km/s
+            self.pert_v_r_ann   *= 1e-5
+            self.pert_v_phi_ann *= 1e-5
+
+            # save annulus grid
+            self.r_ann   = r_ann
+            self.phi_ann = phi_ann
+            self.R_ann   = R_ann
+            self.PHI_ann = PHI_ann
 
-        # linear perturbations read in grid
-        x = self.X[0,:]
-        y = self.Y[:,0]
 
-        # cut square box in linear regime
-        x_cut = x[np.argmin(x < -x_box_size_l) : np.argmin(x < x_box_size_r) + 1]
-        y_cut = y[np.argmin(y < -y_box_size_b) : np.argmin(y < y_box_size_t) + 1]
-
-        # annulus segment grid, granularity from square box
-        r = np.linspace(
-            self.p.r_planet - x_box_size_l*self.p.l, 
-            self.p.r_planet + x_box_size_r*self.p.l, 
-            len(x_cut)
-        )
-        y_ = np.linspace(
-            -y_box_size_b*self.p.l,
-            y_box_size_t*self.p.l,
-            len(y_cut)
-        )
-
-        R, Y_ = np.meshgrid(r, y_)
-
-        # for the points on our annulus segment, use transformations to find the x,y values on the original perturbations grid
-        # if the user has chosen a linear box with very large angular extent, R and Y will be equal at some point, creating an error
-        with np.errstate(all='raise'):
-            try:
 
-                # pick x values for the box
-                if self.p.box_warp:
-                    X_pert_grid = (R - self.p.r_planet) / self.p.l
-                else:
-                    X_pert_grid = (np.sqrt(R**2 - Y_**2) - self.p.r_planet) / self.p.l
 
-                # get phi values
-                PHI = np.arctan2(Y_, np.sqrt(R**2 - Y_**2))
 
-            except FloatingPointError:
-                print("Error: Reduce 'scale_box_ang' parameter.")
-                print("Exiting")
-                sys.exit(1)
-        
-        # unscale Y
-        Y_pert_grid = Y_ / self.p.l
 
-        # cut big perturbations grid to just outside annulus
-        self.r_min = self.p.r_planet - x_box_size_l*self.p.l
-        self.r_max = self.p.r_planet + x_box_size_r*self.p.l
-
-        x_min_global = np.sqrt(self.r_min**2 - (np.amax([y_box_size_t,y_box_size_b])*self.p.l)**2)
-        x_min_local  = (x_min_global - self.p.r_planet) / self.p.l
-
-        # find cut indicies (remember we need to scale back to units of Hill radius )
-        x_cut_i1 = np.argmin(x <  x_min_local)
-        x_cut_i2 = np.argmin(x <  x_box_size_r) + 1
-        y_cut_i1 = np.argmin(y < -y_box_size_b)
-        y_cut_i2 = np.argmin(y <  y_box_size_t) + 1
-
-        # cut grid
-        x_int_cut = x[x_cut_i1 : x_cut_i2]
-        y_int_cut = y[y_cut_i1 : y_cut_i2]
-
-        # cut perturbation arrays
-        cut_v_r   = self.pert_v_r   [y_cut_i1:y_cut_i2, x_cut_i1:x_cut_i2]
-        cut_v_phi = self.pert_v_phi [y_cut_i1:y_cut_i2, x_cut_i1:x_cut_i2]
-        cut_rho   = self.pert_rho   [y_cut_i1:y_cut_i2, x_cut_i1:x_cut_i2]
-
-        if False:
-            plt.contourf(x_int_cut, y_int_cut, cut_rho, cmap="RdBu", vmin=-1, vmax=1, levels=100)
-            plt.show()
-
-        # account for rotation direction
-        if self.p.a_cw == -1:
-            cut_v_r   =  np.flipud(cut_v_r)
-            cut_v_phi = -np.flipud(cut_v_phi)
-            cut_rho   =  np.flipud(cut_rho)
-
-        # interpolation over cut (cartesian) grid
-        interp_v_r   = RectBivariateSpline(y_int_cut, x_int_cut, cut_v_r)
-        interp_v_phi = RectBivariateSpline(y_int_cut, x_int_cut, cut_v_phi)
-        interp_v_rho = RectBivariateSpline(y_int_cut, x_int_cut, cut_rho)
-
-        # evaluate interpolation over our annulus segment
-        self.pert_v_r_ann   = interp_v_r.ev  (Y_pert_grid, X_pert_grid)
-        self.pert_v_phi_ann = interp_v_phi.ev(Y_pert_grid, X_pert_grid)
-        self.pert_rho_ann   = interp_v_rho.ev(Y_pert_grid, X_pert_grid)
-
-        if False:
-            plt.imshow(self.pert_rho_ann, cmap="RdBu", vmin=-1, vmax=1)
-            plt.show()
-
-        # scale to cgs units
-        self.pert_v_r_ann   *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
-        self.pert_v_phi_ann *= self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal)
-        self.pert_rho_ann   *=                   (self.p.m_planet/self.p.m_thermal)
-
-        # scale velocities to km/s
-        self.pert_v_r_ann   *= 1e-5
-        self.pert_v_phi_ann *= 1e-5
-
-        # save annulus grid
-        self.r_ann   = r
-        self.y_ann   = y_
-        self.R_ann   = R
-        self.PHI_ann = PHI
-
-        if False:
-            # plotting (for debugging)
-            _, ax = plt.subplots(subplot_kw=dict(projection='polar'))
-            myplot = ax.contourf(PHI, R, self.pert_rho_ann/(self.p.m_planet/self.p.m_thermal), levels=300, vmin=-1, vmax=1, cmap='RdBu')
-            ax.set_ylim(0, self.p.r_outer)
-            plt.colorbar(myplot)
-            plt.show()
diff --git a/src/wakeflow/model_setup.py b/src/wakeflow/model_setup.py
index db236aa..14bf8cb 100644
--- a/src/wakeflow/model_setup.py
+++ b/src/wakeflow/model_setup.py
@@ -116,17 +116,23 @@ def __init__(self, config: dict) -> None:
         self.malpha = float(config["damping_malpha"])
 
         # numerical parameters
-        self.CFL             = float(config["CFL"])
-        self.smooth_box      = float(config["smooth_box"])
-        self.scale_box_l     = float(config["scale_box_l"])
-        self.scale_box_r     = float(config["scale_box_r"])
-        self.scale_box_ang_t = float(config["scale_box_ang_t"])
-        self.scale_box_ang_b = float(config["scale_box_ang_b"])
-        self.box_warp        = bool (config["box_warp"])
-        self.use_box_IC      = bool (config["use_box_IC"])
-        self.tf_fac          = float(config["tf_fac"])
-        self.use_old_vel     = bool (config["use_old_vel"])
-
+        self.CFL                  = float(config["CFL"])
+        self.smooth_box           = float(config["smooth_box"])
+        self.scale_box_left       = float(config["scale_box_left"])
+        self.scale_box_right      = float(config["scale_box_right"])
+        self.scale_box_ang_top    = float(config["scale_box_ang_top"])
+        self.scale_box_ang_bottom = float(config["scale_box_ang_bottom"])
+        self.box_warp             = bool (config["box_warp"])
+        self.use_box_IC           = bool (config["use_box_IC"])
+        self.tf_fac               = float(config["tf_fac"])
+        self.rot_interp           = bool (config["rot_interp"])
+        self.r_cut_inner_fac      = float(config["r_cut_inner_fac"])
+        
+        # Choice of physics
+        self.use_old_vel          = bool (config["use_old_vel"])
+        self.lin_type             = str  (config["lin_type"])
+        self.nl_wake              = bool (config["nl_wake"])
+        
         # get flaring at r_planet
         self.hr_planet = self.hr * (self.r_planet / self.r_ref) ** (0.5 - self.q)
 
@@ -206,15 +212,15 @@ def _do_sanity_checks(self) -> bool:
                 raise Exception("You must choose grid_type='cylindrical' to use smooth_box=True")
 
         # check linear box scale factor
-        if self.scale_box_l != 1 or self.scale_box_r != 1:
+        if self.scale_box_left != 1 or self.scale_box_right != 1:
             print("WARNING: Changing linear box scale factor can cause strange results.")
 
         # check linear box scale factor
-        if self.scale_box_l != self.scale_box_r:
+        if self.scale_box_left != self.scale_box_right:
             print("WARNING: Using a different linear box scale factor for left and right edge can cause strange results.")
 
         # check linear box scale factor
-        if self.scale_box_ang_t != self.scale_box_ang_b:
+        if self.scale_box_ang_top != self.scale_box_ang_bottom:
             print("WARNING: Using a different linear box scale factor for top and bottom edge can cause strange results.")
 
         # check CFL
@@ -238,8 +244,16 @@ def _do_sanity_checks(self) -> bool:
             print("WARNING: Choosing use_old_vel=True may cause a different velocity output.")
         print("Parameters Ok - continuing")
         return True
+            
+        # check linear perturbations input from user    
+        if self.lin_type != "shearing_sheet" and self.lin_type != "global" and self.lin_type != "simulation":
+            raise Exception("Invalid linear perturbation type. Choose either global or simulation or shearing_sheet)")
+            
+        # Warning on shearing sheet approximation
+        if self.lin_type == "shearing_sheet":
+            print("WARNING: The shearing sheet approximation may be invalid for a given choice of parameters. This may lead to incorrect results")
 
-# read in .yaml file, check keys correspond to parameters and return dictionary of parameters
+ # read in .yaml file, check keys correspond to parameters and return dictionary of parameters
 def _load_config_file(config_file: str, default_config_dict: dict = None) -> dict:
     """Reads .yaml parameter file into a dictionary so that Wakeflow may parse it.
     """
diff --git a/src/wakeflow/non_linear_perts.py b/src/wakeflow/non_linear_perts.py
index 1fdef55..eff881e 100644
--- a/src/wakeflow/non_linear_perts.py
+++ b/src/wakeflow/non_linear_perts.py
@@ -15,7 +15,7 @@
 from copy                   import copy
 from tqdm                   import tqdm
 from .burgers               import _solve_burgers
-from .transformations       import _Eta, _t, _get_chi, _get_dens_vel, _plot_r_t, _t_vector, _get_chi_vector
+from .transformations       import _Eta, _Eta_vector, _t, _get_chi, _get_dens_vel, _plot_r_t, _t_vector, _get_chi_vector, _g, _Lambda_fu, _Lambda_fv
 
 if TYPE_CHECKING:
     from .model_setup       import _Parameters
@@ -72,6 +72,14 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
         y_max = 30
         y_rest = y[(y > -y_max) & (y < y_max)]
 
+        # plot for debugging
+        if False:
+            plt.contourf(x[index_inner:index_outer], y_rest, lp.pert_rho[(y > -y_max) & (y < y_max),index_inner:index_outer], cmap="RdBu", vmin=-1, vmax=1, levels=100)
+            plt.show()
+            plt.contourf(x[index_inner:index_outer], y_rest, lp.pert_rho[(y > -y_max) & (y < y_max),index_inner:index_outer], cmap="RdBu", vmin=-1, vmax=1, levels=100)
+            plt.ylim(-2.0,2.0)
+            plt.show()
+            
         # restrict inner and outer IC
         profile_rest_outer = profile_outer[(y > -y_max) & (y < y_max)]
         profile_rest_inner = profile_inner[(y > -y_max) & (y < y_max)]
@@ -100,24 +108,38 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
         eta_IC_inner = np.zeros(len(y_rest))
         t_IC_outer   = np.zeros(len(y_rest))
         t_IC_inner   = np.zeros(len(y_rest))
+        
+        eta_IC_outer = _Eta_vector(r_IC_outer, phi_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+        eta_IC_inner = _Eta_vector(r_IC_inner, phi_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
 
+        # t for outer and inner
+        t_IC_outer = _t_vector(r_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
+        t_IC_inner = _t_vector(r_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
+
+
+        """
         # perform transformation
         for i in range(len(y_rest)):
 
             # eta for outer and inner
-            eta_IC_outer[i] = _Eta(r_IC_outer[i], phi_IC_outer[i], self.p.r_planet, self.p.hr_planet, self.p.q, -1)
-            eta_IC_inner[i] = _Eta(r_IC_inner[i], phi_IC_inner[i], self.p.r_planet, self.p.hr_planet, self.p.q, -1)
+            eta_IC_outer[i] = _Eta(r_IC_outer[i], phi_IC_outer[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+            eta_IC_inner[i] = _Eta(r_IC_inner[i], phi_IC_inner[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
 
             # t for outer and inner
             t_IC_outer[i] = _t(r_IC_outer[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p)
             t_IC_inner[i] = _t(r_IC_inner[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p)
-
+        """
         # restrict eta range using eta_max
         self.eta_outer = eta_IC_outer[(eta_IC_outer > -eta_max) & (eta_IC_outer < eta_max)]
         self.eta_inner = eta_IC_inner[(eta_IC_inner > -eta_max) & (eta_IC_inner < eta_max)]
         self.profile_outer = profile_rest_outer[(eta_IC_outer > -eta_max) & (eta_IC_outer < eta_max)]
         self.profile_inner = profile_rest_inner[(eta_IC_inner > -eta_max) & (eta_IC_inner < eta_max)]
 
+        if False:
+            plt.plot(self.eta_outer, self.profile_outer)
+            plt.plot(self.eta_inner, self.profile_inner)
+            plt.show()
+
         # set t0
         self.t0_outer = t_IC_outer[0]
         self.t0_inner = t_IC_inner[0]
@@ -126,9 +148,9 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
         if False:
             y_match = -np.sign(lp.x_box) * 0.5 * lp.x_box**2
             y_cut = y[(y - y_match > -eta_max) & (y - y_match < eta_max)]
-            self.profile = profile[(y - y_match > -eta_max) & (y - y_match < eta_max)]
+            self.profile = profile_outer[(y - y_match > -eta_max) & (y - y_match < eta_max)]
             self.eta = y_cut - y_match*np.ones(len(y_cut))
-            self.t0 = _t(self.p.r_planet + self.p.l * lp.x_box, self.p.r_planet, self.p.hr, self.p.q, self.p.p)
+            self.t0 = _t(self.p.r_planet + self.p.l * lp.x_box, self.p.r_planet, self.p.hr, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
 
         # set eta_tilde for outer wake:
         for i in range(len(self.eta_outer)):
@@ -136,36 +158,37 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
                 zero_outer = self.eta_outer[i]
             elif i!= (len(self.eta_outer) - 1) and self.profile_outer[i] * self.profile_outer[i + 1] < 0 and self.eta_outer[i] > -10 and self.eta_outer[i] < 0:
                 zero_outer = 0.5 * (self.eta_outer[i] + self.eta_outer[i + 1])
-        self.eta_tilde_outer = -zero_outer
+        self.eta_tilde_outer = zero_outer
 
         # set eta_tilde for inner wake:
         for i in range(len(self.eta_inner)):
-            if self.profile_inner[i] == 0 and self.eta_inner[i] > -10 and self.eta_inner[i] < 0:
+            if self.profile_inner[i] == 0 and self.eta_inner[i] > 0 and self.eta_inner[i] < 10:
                 zero_inner = self.eta_inner[i]
-            elif i!= (len(self.eta_inner) - 1) and self.profile_inner[i] * self.profile_inner[i + 1] < 0 and self.eta_inner[i] > -10 and self.eta_inner[i] < 0:
+            elif i!= (len(self.eta_inner) - 1) and self.profile_inner[i] * self.profile_inner[i + 1] < 0 and self.eta_inner[i] > 0 and self.eta_inner[i] < 10:
                 zero_inner = 0.5 * (self.eta_inner[i] + self.eta_inner[i + 1])
-        self.eta_tilde_inner = -zero_inner
+        self.eta_tilde_inner = zero_inner
 
         # set C for outer wake:
         deta_outer = self.eta_outer[1] - self.eta_outer[0]
-        profile0_outer = self.profile_outer[self.eta_outer < -self.eta_tilde_outer]
-        C0_outer = -np.trapz(profile0_outer, dx = deta_outer)
-        self.C_outer = (self.p.gamma + 1) * (self.p.m_planet / self.p.m_thermal) * C0_outer / 2**(3/4)
+        profile0_outer = self.profile_outer[self.eta_outer > self.eta_tilde_outer]
+        C0_outer = np.abs(np.trapz(profile0_outer, dx = deta_outer))
+        self.C_outer = (self.p.gamma + 1) * C0_outer / 2**(3/4) #* (self.p.m_planet / self.p.m_thermal)
 
         # set C for inner wake:
         deta_inner = self.eta_inner[1] - self.eta_inner[0]
-        profile0_inner = self.profile_inner[self.eta_inner < -self.eta_tilde_inner]
-        C0_inner = -np.trapz(profile0_inner, dx = deta_inner)
-        self.C_inner = (self.p.gamma + 1) * (self.p.m_planet / self.p.m_thermal) * C0_inner / 2**(3/4)
-
-        #print('     Outer Wake:')
-        #print('         eta_tilde = ', self.eta_tilde_outer)
-        #print('         C0 = ', C0_outer)
-        #print('         t0 = ', self.t0_outer)
-        #print('     Inner Wake:')
-        #print('         eta_tilde = ', self.eta_tilde_inner)
-        #print('         C0 = ', C0_inner)
-        #print('         t0 = ', self.t0_inner)
+        profile0_inner = self.profile_inner[self.eta_inner < self.eta_tilde_inner]
+        C0_inner = np.abs(np.trapz(profile0_inner, dx = deta_inner))
+        self.C_inner = (self.p.gamma + 1) * C0_inner / 2**(3/4) #* (self.p.m_planet / self.p.m_thermal)
+
+        if False:
+            print('     Outer Wake:')
+            print('         eta_tilde = ', self.eta_tilde_outer)
+            print('         C0 = ', C0_outer)
+            print('         t0 = ', self.t0_outer)
+            print('     Inner Wake:')
+            print('         eta_tilde = ', self.eta_tilde_inner)
+            print('         C0 = ', C0_inner)
+            print('         t0 = ', self.t0_inner)
 
         # ======================================
         # === put linear solution into t,eta ===
@@ -176,87 +199,87 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
         beta_p = self.p.m_planet / self.p.m_thermal
         gamma  = self.p.gamma
 
-#        if self.p.show_teta_debug_plots:
-#
-#            index = index_outer
-#
-#            # run square box cut
-#            lp._cut_box_square()
-#
-#            # plot r, t
-#            _plot_r_t(self.p)
-#
-#            # Local Cartesian grid which the linear solution was calculated on, meshgrid version
-#            X, Y = np.meshgrid(lp.x_cut,lp.y_cut)
-#
-#            print("Y MIN MAX", np.min(Y),np.max(Y))
-#
-#            # grab positive x side of linear solution
-#            x_len = len(lp.x_cut) // 2
-#
-#            # restrict box to area considered
-#            X_rest = X[:,x_len:index]
-#            Y_rest = Y[:,x_len:index]
-#
-#            print("Y MIN MAX", np.min(Y_rest),np.max(Y_rest))
-#
-#            # Find chi for linear solution, only taking positive values of x
-#            # The "index" value also makes it so we do not go outside the linear box
-#            linear_profile = lp.cut_rho[:,x_len:index] / np.sqrt(np.abs(X_rest)) ###
-#            chi0_rho = linear_profile[:,-1]  * (gamma+1) * beta_p / 2**(3/4) 
-#
-#            # extract chi from velocity perturbations
-#            v_unit = (self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal))
-#            Rp     = self.p.r_planet
-#            cw     = -self.p.a_cw
-#            hr     = self.p.hr_planet
-#            q      = self.p.q
-#            gamma  = self.p.gamma
-#            csp    = self.p.c_s_planet
-#            p      = self.p.p
-#            l      = self.p.l
-#            linear_profile_v_r   = v_unit * (lp.cut_v_r[:,x_len:index]   / np.sqrt(np.abs(X_rest))) ###
-#            linear_profile_v_phi = v_unit * (lp.cut_v_phi[:,x_len:index] / np.sqrt(np.abs(X_rest))) ###
-#            R0  = np.sqrt((Rp + l*X_rest[:,-1])**2 + (l*Y_rest[:,-1])**2)
-#            Lfu = _Lambda_fu(R0, Rp, csp, hr, gamma, q, p)
-#            Lfv = _Lambda_fv(R0, Rp, csp, hr, gamma, q, p)
-#            chi0_v_r = linear_profile_v_r[:,-1] / (np.sign(R0 - Rp) * Lfu)
-#            chi0_v_phi = linear_profile_v_phi[:,-1] / (np.sign(R0 - Rp) * Lfv * (-cw))
-#
-#            # Find etas for the box using IC method
-#            linear_eta = Y_rest + np.sign(X_rest) * 0.5 * X_rest**2
-#
-#            # Find etas using proper transformations
-#            #hp = self.p.hr * self.p.r_planet
-#            hp = self.p.hr_planet * self.p.r_planet
-#            x_glob = 2 * hp * X_rest / 3.
-#            y_glob = 2 * hp * Y_rest / 3.
-#            r_glob = x_glob + self.p.r_planet
-#            phi_glob = y_glob / self.p.r_planet
-#            eta_lin = np.zeros(linear_profile.shape)
-#            t_lin = np.zeros(linear_profile.shape)
-#            for i in range(eta_lin.shape[0]):
-#                #print(str(i))
-#                for j in range(eta_lin.shape[1]):
-#                    eta_lin[i,j] = _Eta(r_glob[i,j], phi_glob[i,j], self.p.r_planet, self.p.hr_planet, self.p.q, -self.p.a_cw)
-#                    t_lin[i,j] = _t(r_glob[i,j], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p)
-#
-#            # Plot Chi vs eta when using eta transformation used for IC cut out
-#
-#            #for i in range(0, len(linear_eta[0,:]), 10):
-#            #    plt.plot(linear_eta[:,i], linear_profile[:,i])
-#            #    plt.plot(eta_lin[:,i], linear_profile[:,i])
-#            #    plt.show()
-#
-#            plt.plot(eta_lin[:,-1], chi0_rho, label="chi rho")
-#            plt.plot(eta_lin[:,-1], chi0_v_r, label="chi vr")
-#            plt.plot(eta_lin[:,-1], chi0_v_phi, label="chi vphi")
-#            plt.legend(loc="lower left")
-#            plt.xlim(-10,10)
-#            plt.xlabel('$\eta$')
-#            plt.ylabel('$\chi$')
-#            plt.show()
-#
+        if False:
+
+            index = index_outer
+
+            # run square box cut
+            lp._cut_box_square()
+
+            # plot r, t
+            _plot_r_t(self.p)
+
+            # Local Cartesian grid which the linear solution was calculated on, meshgrid version
+            X, Y = np.meshgrid(lp.x_cut,lp.y_cut)
+
+            print("Y MIN MAX", np.min(Y),np.max(Y))
+
+            # grab positive x side of linear solution
+            x_len = len(lp.x_cut) // 2
+
+            # restrict box to area considered
+            X_rest = X[:,x_len:index]
+            Y_rest = Y[:,x_len:index]
+
+            print("Y MIN MAX", np.min(Y_rest),np.max(Y_rest))
+
+            # Find chi for linear solution, only taking positive values of x
+            # The "index" value also makes it so we do not go outside the linear box
+            linear_profile = lp.cut_rho[:,x_len:index] / np.sqrt(np.abs(X_rest)) ###
+            chi0_rho = linear_profile[:,-1]  * (gamma+1) * beta_p / 2**(3/4) 
+
+            # extract chi from velocity perturbations
+            v_unit = (self.p.c_s_planet*(self.p.m_planet/self.p.m_thermal))
+            Rp     = self.p.r_planet
+            cw     = -self.p.a_cw
+            hr     = self.p.hr_planet
+            q      = self.p.q
+            gamma  = self.p.gamma
+            csp    = self.p.c_s_planet
+            p      = self.p.p
+            l      = self.p.l
+            linear_profile_v_r   = v_unit * (lp.cut_v_r[:,x_len:index]   / np.sqrt(np.abs(X_rest))) ###
+            linear_profile_v_phi = v_unit * (lp.cut_v_phi[:,x_len:index] / np.sqrt(np.abs(X_rest))) ###
+            R0  = np.sqrt((Rp + l*X_rest[:,-1])**2 + (l*Y_rest[:,-1])**2)
+            Lfu = _Lambda_fu(R0, Rp, csp, hr, gamma, q, p)
+            Lfv = _Lambda_fv(R0, Rp, csp, hr, gamma, q, p)
+            chi0_v_r = linear_profile_v_r[:,-1] / (np.sign(R0 - Rp) * Lfu)
+            chi0_v_phi = linear_profile_v_phi[:,-1] / (np.sign(R0 - Rp) * Lfv * (-cw))
+
+            # Find etas for the box using IC method
+            linear_eta = Y_rest + np.sign(X_rest) * 0.5 * X_rest**2
+
+            # Find etas using proper transformations
+            #hp = self.p.hr * self.p.r_planet
+            hp = self.p.hr_planet * self.p.r_planet
+            x_glob = 2 * hp * X_rest / 3.
+            y_glob = 2 * hp * Y_rest / 3.
+            r_glob = x_glob + self.p.r_planet
+            phi_glob = y_glob / self.p.r_planet
+            eta_lin = np.zeros(linear_profile.shape)
+            t_lin = np.zeros(linear_profile.shape)
+            for i in range(eta_lin.shape[0]):
+                #print(str(i))
+                for j in range(eta_lin.shape[1]):
+                    eta_lin[i,j] = _Eta(r_glob[i,j], phi_glob[i,j], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -self.p.a_cw, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+                    t_lin[i,j] = _t(r_glob[i,j], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
+
+            # Plot Chi vs eta when using eta transformation used for IC cut out
+
+            #for i in range(0, len(linear_eta[0,:]), 10):
+            #    plt.plot(linear_eta[:,i], linear_profile[:,i])
+            #    plt.plot(eta_lin[:,i], linear_profile[:,i])
+            #    plt.show()
+
+            plt.plot(eta_lin[:,-1], chi0_rho, label="chi rho")
+            plt.plot(eta_lin[:,-1], chi0_v_r, label="chi vr")
+            plt.plot(eta_lin[:,-1], chi0_v_phi, label="chi vphi")
+            plt.legend(loc="lower left")
+            plt.xlim(-10,10)
+            plt.xlabel(r'$\eta$')
+            plt.ylabel(r'$\chi$')
+            plt.show()
+
 #            if True: # save results to file for plotting
 #
 #                name = "sim_further"
@@ -265,58 +288,58 @@ def _extract_ICs(self, LinearPerts: '_LinearPerts') -> None:
 #                np.save(f"{name}_eta.npy", eta_save)
 #                np.save(f"{name}_chi.npy", chi_save)
 #
-#            plt.plot(linear_eta[:,-1], linear_profile[:,-1], label="Approximate $\eta$ transformation")
-#            plt.plot(eta_lin[:,-1], linear_profile[:,-1], label="Full $\eta$ transformation")
-#            plt.plot(self.eta, self.profile, ls='--', label="Actual IC used")
-#            plt.legend(loc="lower left")
-#            plt.xlim(-10,10)
-#            plt.xlabel('$\eta$')
-#            plt.ylabel('$\chi$')
-#            plt.show()
-#
-#            plt.scatter(t_lin, eta_lin)
-#            plt.show()
-#
-#            #eta_lin = linear_eta    ### TURN THIS ON AND OFF TO CHANGE ETA TRANSFORMATION --> COMMENT OUT TO USE INTEGRAL TRANSFORM
-#
-#            # get grid regular in (t,eta) to interpolate onto
-#            dt_lin = np.diff(np.sort(t_lin.flatten())).mean()       # find time step interval to use on regular grid from t_lin
-#                                                                    # takes minimum separation between points in original t_lin grid
-#            t_lin_reg = np.arange(0, self.t0, dt_lin)
-#            eta_lin_reg = copy(self.eta)
-#            T_lin_reg, ETA_lin_reg = np.meshgrid(t_lin_reg, eta_lin_reg)
-#
-#            # interpolate solution over irregular grid onto regular grid
-#            print("Interpolating Linear solution into (t, eta) space")
-#            lin_solution = griddata(
-#                (t_lin.flatten(), eta_lin.flatten()), 
-#                linear_profile.flatten(),
-#                (T_lin_reg, ETA_lin_reg),
-#                method='linear'
-#            )
-#
-#            #plt.plot(linear_eta[:,-1], linear_profile[:,-1])       # approx
-#            plt.plot(eta_lin[:,-1], linear_profile[:,-1])           # integral trans
-#            plt.plot(eta_lin_reg, lin_solution[:,-1])               # interpolated integral trans
-#            #plt.plot(self.eta, self.profile, c='k')                # IC
-#            plt.show()
-#
-#            # plotting (for debugging)
-#            _, ax = plt.subplots()
-#            myplot = ax.contourf(t_lin_reg, eta_lin_reg, np.nan_to_num(lin_solution), levels=np.arange(-4,4,0.1), cmap='RdBu', extend='both')
-#            plt.colorbar(myplot)
-#            plt.show()
-#
-#            # stick this into the solutions array, update later code so that solutions array builds on this one
-#            self.linear_solution = np.nan_to_num(lin_solution)
-#            self.linear_t = t_lin_reg
-#
-#        else:
+            plt.plot(linear_eta[:,-1], linear_profile[:,-1], label=r"Approximate $\eta$ transformation")
+            plt.plot(eta_lin[:,-1], linear_profile[:,-1], label=r"Full $\eta$ transformation")
+            #plt.plot(self.eta, self.profile, ls='--', label="Actual IC used")
+            plt.legend(loc="lower left")
+            plt.xlim(-10,10)
+            plt.xlabel(r'$\eta$')
+            plt.ylabel(r'$\chi$')
+            plt.show()
+
+            plt.scatter(t_lin, eta_lin)
+            plt.show()
+
+            #eta_lin = linear_eta    ### TURN THIS ON AND OFF TO CHANGE ETA TRANSFORMATION --> COMMENT OUT TO USE INTEGRAL TRANSFORM
+
+            # get grid regular in (t,eta) to interpolate onto
+            dt_lin = np.diff(np.sort(t_lin.flatten())).mean()       # find time step interval to use on regular grid from t_lin
+                                                                    # takes minimum separation between points in original t_lin grid
+            t_lin_reg = np.arange(0, self.t0, dt_lin)
+            eta_lin_reg = copy(self.eta)
+            T_lin_reg, ETA_lin_reg = np.meshgrid(t_lin_reg, eta_lin_reg)
+
+            # interpolate solution over irregular grid onto regular grid
+            print("Interpolating Linear solution into (t, eta) space")
+            lin_solution = griddata(
+                (t_lin.flatten(), eta_lin.flatten()), 
+                linear_profile.flatten(),
+                (T_lin_reg, ETA_lin_reg),
+                method='linear'
+            )
+
+            #plt.plot(linear_eta[:,-1], linear_profile[:,-1])       # approx
+            plt.plot(eta_lin[:,-1], linear_profile[:,-1])           # integral trans
+            plt.plot(eta_lin_reg, lin_solution[:,-1])               # interpolated integral trans
+            #plt.plot(self.eta, self.profile, c='k')                # IC
+            plt.show()
+
+            # plotting (for debugging)
+            _, ax = plt.subplots()
+            myplot = ax.contourf(t_lin_reg, eta_lin_reg, np.nan_to_num(lin_solution), levels=np.arange(-4,4,0.1), cmap='RdBu', extend='both')
+            plt.colorbar(myplot)
+            plt.show()
+
+            # stick this into the solutions array, update later code so that solutions array builds on this one
+            self.linear_solution = np.nan_to_num(lin_solution)
+            self.linear_t = t_lin_reg
+
+        else:
             
-        self.linear_solution = 0
-        self.linear_t = 0
+            self.linear_solution = 0
+            self.linear_t = 0
 
-    # alternative IC extraction using edge of box
+    # alternative IC extraction using an annulus instead of the linear square box
     def _extract_ICs_ann(self, LinearPerts: '_LinearPerts') -> None:
         """Alternate initial condition extraction where the IC is read from the edges of the box as included in the final
         solution. Usually, far more y-extent of the linear regime is used than this. Using this method will invalidate the
@@ -335,12 +358,9 @@ def _extract_ICs_ann(self, LinearPerts: '_LinearPerts') -> None:
         r_IC_outer   = lp.r_ann[-1]
         r_IC_inner   = lp.r_ann[ 0]
 
-        # edge of box in local coords
-        x_box_outer = (r_IC_outer - self.p.r_planet) / self.p.l
-        x_box_inner = (r_IC_inner - self.p.r_planet) / self.p.l
-
-        self.profile_outer = (lp.pert_rho_ann[:,-1] / beta_p) / np.sqrt(np.abs(x_box_outer))
-        self.profile_inner = (lp.pert_rho_ann[:, 0] / beta_p) / np.sqrt(np.abs(x_box_inner))
+        #Evaluate Chi profile using global transformation
+        self.profile_outer = (lp.pert_rho_ann[:,-1] / beta_p) * (_g(r_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p) * (self.p.gamma + 1) / 2)
+        self.profile_inner = (lp.pert_rho_ann[:, 0] / beta_p) * (_g(r_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p) * (self.p.gamma + 1) / 2)
 
         if False:
          plt.plot(phi_IC_outer, self.profile_outer, label="outer")
@@ -349,18 +369,35 @@ def _extract_ICs_ann(self, LinearPerts: '_LinearPerts') -> None:
          plt.show()
 
         # find t points
-        t_IC_outer = _t(r_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p)
-        t_IC_inner = _t(r_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p)
+        t_IC_outer = _t(r_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
+        t_IC_inner = _t(r_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
 
         # initialise arrays for corresponding eta points
         self.eta_outer = np.zeros(len(phi_IC_outer))
         self.eta_inner = np.zeros(len(phi_IC_outer))
 
         # perform transformation
+        self.eta_outer = _Eta_vector(r_IC_outer, phi_IC_outer, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+        self.eta_inner = _Eta_vector(r_IC_inner, phi_IC_inner, self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+        
+        #Sorting eta and profile in order to have eta in ascending order
+        self.eta_outer = self.eta_outer[:-1]
+        idx_srt = np.argsort(self.eta_outer)
+        self.eta_outer = self.eta_outer[idx_srt]
+        self.profile_outer = self.profile_outer[idx_srt]
+
+        #Sorting eta and profile in order to have eta in ascending order
+        self.eta_inner = self.eta_inner[:-1]
+        idx_srt = np.argsort(self.eta_inner)
+        self.eta_inner = self.eta_inner[idx_srt]
+        self.profile_inner = self.profile_inner[idx_srt]
+        '''
         for i in range(len(phi_IC_outer)):
-            self.eta_outer[i] = _Eta(r_IC_outer, phi_IC_outer[i], self.p.r_planet, self.p.hr_planet, self.p.q, -1)
-            self.eta_inner[i] = _Eta(r_IC_inner, phi_IC_inner[i], self.p.r_planet, self.p.hr_planet, self.p.q, -1)
-
+            self.eta_outer[i] = _Eta(r_IC_outer, phi_IC_outer[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+            self.eta_inner[i] = _Eta(r_IC_inner, phi_IC_inner[i], self.p.r_planet, self.p.hr_planet, self.p.q, self.p.p, -1, self.p.m_planet, self.p.m_thermal, self.p.nl_wake)
+        '''
+        if False:
+            plt.plot(self.eta_outer,self.profile_outer)
         # set t0
         self.t0_outer = t_IC_outer
         self.t0_inner = t_IC_inner
@@ -371,36 +408,38 @@ def _extract_ICs_ann(self, LinearPerts: '_LinearPerts') -> None:
                 zero_outer = self.eta_outer[i]
             elif i!= (len(self.eta_outer) - 1) and self.profile_outer[i] * self.profile_outer[i + 1] < 0 and self.eta_outer[i] > -10 and self.eta_outer[i] < 0:
                 zero_outer = 0.5 * (self.eta_outer[i] + self.eta_outer[i + 1])
-        self.eta_tilde_outer = -zero_outer
+        self.eta_tilde_outer = zero_outer
 
+        if False:
+            plt.plot(self.eta_inner,self.profile_inner)
         # set eta_tilde for inner wake:
         for i in range(len(self.eta_inner)):
-            if self.profile_inner[i] == 0 and self.eta_inner[i] > -10 and self.eta_inner[i] < 0:
+            if self.profile_inner[i] == 0 and self.eta_inner[i] > 0 and self.eta_inner[i] < 10:
                 zero_inner = self.eta_inner[i]
-            elif i!= (len(self.eta_inner) - 1) and self.profile_inner[i] * self.profile_inner[i + 1] < 0 and self.eta_inner[i] > -10 and self.eta_inner[i] < 0:
+            elif i!= (len(self.eta_inner) - 1) and self.profile_inner[i] * self.profile_inner[i + 1] < 0 and self.eta_inner[i] > 0 and self.eta_inner[i] < 10:
                 zero_inner = 0.5 * (self.eta_inner[i] + self.eta_inner[i + 1])
-        self.eta_tilde_inner = -zero_inner
+        self.eta_tilde_inner = zero_inner
 
         # set C for outer wake:
         deta_outer = self.eta_outer[1] - self.eta_outer[0]
-        profile0_outer = self.profile_outer[self.eta_outer < -self.eta_tilde_outer]
-        C0_outer = -np.trapz(profile0_outer, dx = deta_outer)
-        self.C_outer = (self.p.gamma + 1) * (self.p.m_planet / self.p.m_thermal) * C0_outer / 2**(3/4)
+        profile0_outer = self.profile_outer[self.eta_outer > self.eta_tilde_outer]
+        C0_outer = np.abs(np.trapz(profile0_outer, dx = deta_outer))
+        self.C_outer =  C0_outer 
 
         # set C for inner wake:
         deta_inner = self.eta_inner[1] - self.eta_inner[0]
-        profile0_inner = self.profile_inner[self.eta_inner < -self.eta_tilde_inner]
-        C0_inner = -np.trapz(profile0_inner, dx = deta_inner)
-        self.C_inner = (self.p.gamma + 1) * (self.p.m_planet / self.p.m_thermal) * C0_inner / 2**(3/4)
-
-        #print('     Outer Wake:')
-        #print('         eta_tilde = ', self.eta_tilde_outer)
-        #print('         C0 = ', C0_outer)
-        #print('         t0 = ', self.t0_outer)
-        #print('     Inner Wake:')
-        #print('         eta_tilde = ', self.eta_tilde_inner)
-        #print('         C0 = ', C0_inner)
-        #print('         t0 = ', self.t0_inner)
+        profile0_inner = self.profile_inner[self.eta_inner < self.eta_tilde_inner]
+        C0_inner = np.abs(np.trapz(profile0_inner, dx = deta_inner))
+        self.C_inner = C0_inner
+        if False:
+            print('     Outer Wake:')
+            print('         eta_tilde = ', self.eta_tilde_outer)
+            print('         C0 = ', C0_outer)
+            print('         t0 = ', self.t0_outer)
+            print('     Inner Wake:')
+            print('         eta_tilde = ', self.eta_tilde_inner)
+            print('         C0 = ', C0_inner)
+            print('         t0 = ', self.t0_inner)
 
         # hard to explain, but they're not needed so set to zero
         self.linear_solution = 0
@@ -414,6 +453,9 @@ def _get_non_linear_perts(self) -> None:
 
         beta_p = self.p.m_planet / self.p.m_thermal
 
+        # outer edge limit for t. Needed to save computation time in case the disc parameters prevents the wake reaching the asymptotic limit before hitting the outer edge of the disc
+        self.t_edge_outer = _t_vector(np.array(self.p.r_outer), self.p.r_planet, self.p.hr, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
+
         print('Propagating outer wake... ')
 
         timer_0 = time.perf_counter()
@@ -431,6 +473,7 @@ def _get_non_linear_perts(self) -> None:
             self.linear_t, 
             self.p.show_teta_debug_plots,
             self.p.tf_fac,
+            self.t_edge_outer
         )
 
         timer_1 = time.perf_counter()
@@ -440,6 +483,9 @@ def _get_non_linear_perts(self) -> None:
         print('Propagating inner wake... ')
 
         timer_0 = time.perf_counter()
+        
+        # inner edge limit for t. Needed to save computation time as we usually can not resolve the inner disc in observations. The r_cut_inner_fac can be changed to increase or decrease the computation in the inner disc, I still need to find a better way to do this
+        self.t_edge_inner = _t_vector(np.array(self.p.r_planet/self.p.r_cut_inner_fac), self.p.r_planet, self.p.hr, self.p.q, self.p.p, self.p.m_planet, self.p.m_thermal)
 
         time_inner, eta_inner, solution_inner = _solve_burgers(
             self.eta_inner, 
@@ -453,7 +499,8 @@ def _get_non_linear_perts(self) -> None:
             self.linear_solution, 
             self.linear_t, 
             self.p.show_teta_debug_plots,
-            self.p.tf_fac
+            self.p.tf_fac,
+            self.t_edge_inner
         )
 
         timer_1 = time.perf_counter()
@@ -464,6 +511,8 @@ def _get_non_linear_perts(self) -> None:
         tf_outer = time_outer[-1]
         tf_inner = time_inner[-1]
 
+       
+
         # needed constants from solutions, inner and outer
         eta_tilde_outer = self.eta_tilde_outer
         eta_tilde_inner = self.eta_tilde_inner
@@ -474,14 +523,17 @@ def _get_non_linear_perts(self) -> None:
 
         # parameters of run
         Rp      = self.p.r_planet
-        x_match_l = 2*self.p.scale_box_l
-        x_match_r = 2*self.p.scale_box_r
+        x_match_l = 2*self.p.scale_box_left
+        x_match_r = 2*self.p.scale_box_right
         l       = self.p.l
         cw      = -self.p.a_cw
         hr      = self.p.hr_planet
         q       = self.p.q
         csp     = self.p.c_s_planet
         p       = self.p.p
+        m_p     = self.p.m_planet
+        m_th    = self.p.m_thermal
+        nl_wake = self.p.nl_wake
 
         # physical parameters
         gamma = self.p.gamma
@@ -501,7 +553,11 @@ def _get_non_linear_perts(self) -> None:
             r_grid = np.sqrt(x_grid**2 + y_grid**2)
             pphi_grid = np.arctan2(y_grid, x_grid) - self.p.phi_planet
 
-            tt = _t_vector(r_grid, Rp, hr, q, p)
+            tt = _t_vector(r_grid, Rp, hr, q, p, m_p, m_th)
+            if False:
+                plt.plot(eta_outer)
+                plt.plot(t0_outer + time_outer)
+                plt.show()
 
             CChi = _get_chi_vector(
                 pphi_grid, 
@@ -529,6 +585,9 @@ def _get_non_linear_perts(self) -> None:
                 hr, 
                 q, 
                 p,
+                m_p,
+                m_th,
+                nl_wake
             )
 
             # COMPUTE DENSITY AND VELOCITY PERTURBATIONS
@@ -542,7 +601,9 @@ def _get_non_linear_perts(self) -> None:
                 hr, 
                 q, 
                 p,
-                self.p.use_old_vel
+                self.p.use_old_vel,
+                self.p.m_planet,
+                self.p.m_thermal
             )
 
         # if using a cylindrical grid
@@ -552,7 +613,7 @@ def _get_non_linear_perts(self) -> None:
 
             r_grid, pphi_grid = np.meshgrid(r, phi)
 
-            tt = _t_vector(r_grid, Rp, hr, q, p)
+            tt = _t_vector(r_grid, Rp, hr, q, p, m_p, m_th)
 
             Chi = _get_chi_vector(
                 pphi_grid, 
@@ -579,7 +640,10 @@ def _get_non_linear_perts(self) -> None:
                 cw, 
                 hr, 
                 q, 
-                p
+                p,
+                m_p,
+                m_th, 
+                nl_wake
             )
 
             # COMPUTE DENSITY AND VELOCITY PERTURBATIONS
@@ -593,7 +657,9 @@ def _get_non_linear_perts(self) -> None:
                 hr, 
                 q, 
                 p,
-                self.p.use_old_vel
+                self.p.use_old_vel,
+                self.p.m_planet,
+                self.p.m_thermal
             )
         
         timer_1 = time.perf_counter()
diff --git a/src/wakeflow/transformations.py b/src/wakeflow/transformations.py
index a7feb2e..d4305ae 100644
--- a/src/wakeflow/transformations.py
+++ b/src/wakeflow/transformations.py
@@ -14,19 +14,25 @@
 
 # TODO: update code to be faster and more efficient, most of the code in this file is still leftover from the old version (Analytical Kinks)
 
+
 # wake shape
-def _phi_wake(r, Rp, hr, q, cw): 
-    """Eq. (4) Bollati et al. 2021
+def _phi_wake(r, Rp, hr, q, p, cw, m_p, m_th, nl_wake): 
+    """Eq. (4) Bollati et al. 2021, non-linear correction Eq. (32) Cimerman & Rafikov (2022)
     """
     rr = r / Rp
-    return -cw * np.sign(r - Rp) * (1 / hr) * (rr**(q - 0.5) / (q - 0.5) - rr**(q + 1) / (q + 1) - 3 / ((2 * q - 1) * (q + 1)))
+    phi_l = np.sign(r - Rp) * (1 / hr) * (rr**(q - 0.5) / (q - 0.5) - rr**(q + 1) / (q + 1) - 3 / ((2 * q - 1) * (q + 1)))
+    if nl_wake:
+        phi_nl = np.sign(r - Rp) * hr * np.sqrt(_t_vector(r, Rp, hr, q, p, m_p, m_th))
+    else:
+        phi_nl = 0
+    return -cw * (phi_l + phi_nl)
 
 # eta coordinate transformation
-def _Eta(r, phi, Rp, hr, q, cw):
+def _Eta(r, phi, Rp, hr, q, p, cw, m_p, m_th, nl_wake):
     """Eq. (14) Bollati et al. 2021
     """
     coeff    = 1.5 / hr
-    phi_w    = _mod2pi(_phi_wake(r, Rp, hr, q, cw))
+    phi_w    = _mod2pi(_phi_wake(r, Rp, hr, q, p, cw, m_p, m_th, nl_wake))
     deltaphi = phi - phi_w
 
     if deltaphi > np.pi:
@@ -37,7 +43,7 @@ def _Eta(r, phi, Rp, hr, q, cw):
 
     return coeff * deltaphi
 
-def _Eta_vector(r, phi, Rp, hr, q, cw):
+def _Eta_vector(r, phi, Rp, hr, q, p, cw, m_p, m_th, nl_wake):
     """Eq. (14) Bollati et al. 2021
 
     Vectorised version of _Eta.
@@ -45,7 +51,7 @@ def _Eta_vector(r, phi, Rp, hr, q, cw):
     modulus operators and constant offsets.
     """
     coeff    = 1.5 / hr
-    phi_w    = _phi_wake(r, Rp, hr, q, cw) % (2*np.pi)
+    phi_w    = _phi_wake(r, Rp, hr, q, p, cw, m_p, m_th, nl_wake) % (2*np.pi)
     deltaphi = (phi - phi_w + np.pi) % (2*np.pi) - np.pi
 
     return coeff * deltaphi
@@ -76,7 +82,7 @@ def _t_integrand(x, q, p):
 def _t_integral(up, q, p):
     return  quad(_t_integrand, 1, up, args=(q,p))[0]
 
-def _t_vector(rr, Rp, hr, q, p):
+def _t_vector(rr, Rp, hr, q, p, m_p, m_th):
     """Equation (43) Rafikov 2002    (Eq. 13 Bollati et al. 2021)
     This is a vectorised version of _t.
     _t computes an integral where the integrand is independent of the radius r.
@@ -138,16 +144,16 @@ def _t_vector(rr, Rp, hr, q, p):
     module_integral = np.zeros_like(sorted_results)
     module_integral[integral_bounds_sorted_indices] = np.abs(sorted_results)
 
-    coeff = 3 * hr**(-5 / 2) / (2**(5 / 4))
+    coeff = 3 * hr**(-5 / 2) / (2**(5 / 4)) * (m_p/m_th)
     return np.reshape(coeff * module_integral, shape)
 
 
 # t coordinate transformation
-def _t(r, Rp, hr, q, p):
+def _t(r, Rp, hr, q, p, m_p, m_th):
     """Equation (43) Rafikov 2002    (Eq. 13 Bollati et al. 2021)
     """
     module_integral = np.abs(_t_integral(r / Rp, q, p))
-    coeff = 3 * hr**(-5 / 2) / (2**(5 / 4))
+    coeff = 3 * hr**(-5 / 2) / (2**(5 / 4)) * (m_p/m_th)
     return coeff * module_integral
 
 # g(r) quantity calculation
@@ -202,7 +208,10 @@ def _get_chi(
     hr, 
     q, 
     p,
-    t1
+    t1,
+    m_p,
+    m_th, 
+    nl_wake
 ):
     # COMPUTATION OF Chi
 
@@ -216,7 +225,7 @@ def _get_chi(
     #    print(t1-t1_orig)
 
 
-    eta1 = _Eta(rr, pphi, Rp, hr, q, cw)
+    eta1 = _Eta(rr, pphi, Rp, hr, q, p, cw, m_p, m_th, nl_wake)
 
     # If the point is in the outer disk, use the outer wake solution
     if (rr - Rp) > 0:
@@ -246,7 +255,7 @@ def _get_chi(
             extr_right = +cw * np.sign(rr - Rp) * eta_tilde_outer + np.sqrt(2 * C_outer * (t1 - t0_outer))
 
             if eta1 > extr_left and eta1 < extr_right:
-                Chi = (-cw * np.sign(rr - Rp) * eta1 + eta_tilde_outer) / (t1 - t0_outer)  #eq.(29) nonlinear.pdf
+                Chi = (-cw * np.sign(rr - Rp) * (eta1 - eta_tilde_outer)) / (t1 - t0_outer)  #eq.(29) nonlinear.pdf
             else:
                 Chi = 0
 
@@ -278,7 +287,7 @@ def _get_chi(
             extr_right = +cw * np.sign(rr - Rp) * eta_tilde_inner + np.sqrt(2 * C_inner * (t1 - t0_inner))
 
             if eta1 > extr_left and eta1 < extr_right:
-                Chi = (-cw * np.sign(rr - Rp) * eta1 + eta_tilde_inner) / (t1 - t0_inner)  #eq.(29) nonlinear.pdf
+                Chi = (-cw * np.sign(rr - Rp) * (eta1 - eta_tilde_inner)) / (t1 - t0_inner)  #eq.(29) nonlinear.pdf
             else:
                 Chi = 0
     
@@ -311,6 +320,9 @@ def _get_chi_vector(
     hr, 
     q, 
     p,
+    m_p,
+    m_th,
+    nl_wake
 ): 
     """
     This is a vectorised version of the previous _get_chi function. Uses masks rather than logical statements to replace
@@ -321,7 +333,7 @@ def _get_chi_vector(
 
     Chi = np.zeros_like(rr)
 
-    eta_array = _Eta_vector(rr, pphi, Rp, hr, q, cw)
+    eta_array = _Eta_vector(rr, pphi, Rp, hr, q, p, cw, m_p, m_th, nl_wake)
 
     # Inner and outer masks will account for the annulus directly.
     outer_mask = rr - Rp >= x_match_r * l
@@ -354,7 +366,7 @@ def _get_chi_vector(
         np.logical_and(~before_N_wave_mask,
         np.abs( eta_array - cw * np.sign(rr - Rp) * eta_tilde_outer) <= np.sqrt(2 * C_outer * np.abs(tt - t0_outer))
         ))
-    Chi[m] = (-cw * np.sign(rr[m] - Rp) * eta_array[m] + eta_tilde_outer) / (tt[m] - t0_outer)
+    Chi[m] = (-cw * np.sign(rr[m] - Rp) * (eta_array[m] - eta_tilde_outer)) / (tt[m] - t0_outer)
 
 
     """
@@ -392,7 +404,7 @@ def _get_chi_vector(
         np.logical_and(~before_N_wave_mask,
         np.abs( eta_array - cw * np.sign(rr - Rp) * eta_tilde_inner) < np.sqrt(2 * C_inner * np.abs(tt - t0_inner))
         ))
-    Chi[m] = (-cw * np.sign(rr[m] - Rp) * eta_array[m] + eta_tilde_inner) / (tt[m] - t0_inner)
+    Chi[m] = (-cw * np.sign(rr[m] - Rp) * (eta_array[m] - eta_tilde_inner)) / (tt[m] - t0_inner)
 
     """
     if (rr - Rp) <= -x_match*l:
@@ -412,16 +424,16 @@ def _get_chi_vector(
     return Chi
 
 # get the density and velocity perturbations at the grid point from chi
-def _get_dens_vel(rr, Chi, gamma, Rp, cw, csp, hr, q, p, use_old_vel):
+def _get_dens_vel(rr, Chi, gamma, Rp, cw, csp, hr, q, p, use_old_vel, m_p, m_th):
 
     g1  = _g(rr, Rp, hr, q, p)
-    dnl = Chi * 2 / (g1 * (gamma + 1))     # Eq. (11) Bollati et al. 2021
+    dnl = Chi * 2 / (g1 * (gamma + 1))*(m_p/m_th)     # Eq. (11) Bollati et al. 2021
     
     if use_old_vel == True:
         Lfu = _Lambda_fu(rr, Rp, csp, hr, gamma, q, p)
         Lfv = _Lambda_fv(rr, Rp, csp, hr, gamma, q, p)
-        unl = np.sign(rr - Rp) * Lfu * Chi           # Eq. (23) Bollati et al. 2021
-        vnl = np.sign(rr - Rp) * Lfv * Chi * (-cw) # Eq. (24) Bollati et al. 2021 (the sign of v is reversed if we change cw)
+        unl = np.sign(rr - Rp) * Lfu * Chi*(m_p/m_th)            # Eq. (23) Bollati et al. 2021
+        vnl = np.sign(rr - Rp) * Lfv * Chi * (-cw)*(m_p/m_th)    # Eq. (24) Bollati et al. 2021 (the sign of v is reversed if we change cw)
     else:
         #psi = (np.power(dnl + 1, (gamma-1)/2) - 1) * (gamma+1) / (gamma-1)
         psi = ((gamma+1) / (gamma-1)) * np.sign(dnl + 1) * ((np.abs(dnl + 1)) ** ((gamma - 1) / 2) - 1)
@@ -441,7 +453,7 @@ def _plot_r_t(params):
     r = np.linspace(params.r_planet, params.r_outer, 1000)
     times = []
     for rad in r:
-        times.append(_t(rad, params.r_planet, params.hr_planet, params.q, params.p))
+        times.append(_t(rad, params.r_planet, params.hr_planet, params.q, params.p, params.m_planet, params.m_thermal))
 
     plt.plot(r, times)
     plt.xlabel("r")
diff --git a/src/wakeflow/wakeflow.py b/src/wakeflow/wakeflow.py
index 995cfd8..59dd742 100644
--- a/src/wakeflow/wakeflow.py
+++ b/src/wakeflow/wakeflow.py
@@ -12,6 +12,7 @@
 from .linear_perts        import _LinearPerts
 from .non_linear_perts    import _NonLinearPerts
 #from phantom_interface   import PhantomDump
+from typing import Optional
 
 # class for use by the user to interact with Wakeflow
 class WakeflowModel():
@@ -66,6 +67,7 @@ def configure(
         include_linear:      bool = True,
         save_perturbations:  bool = True,
         save_total:          bool = True,
+        rot_interp:          bool = False,
         write_FITS:          bool = False,
         run_mcfost:          bool = False,
         inclination:        float = -225,
@@ -141,6 +143,8 @@ def configure(
             Save the perturbations?
         save_total : bool
             Save the totals (perturbations + background disk)?
+        rot_interp : bool
+            Extend grid by a factor sqrt(2) to account for interpolation on corners over rotated grid (working only for cartesian)
         write_FITS : bool
             Generate a .FITS file to run in MCFOST? Requires "mcfost" grid type.
         run_mcfost : bool
@@ -167,15 +171,19 @@ def configure(
         damping_malpha        = 0.0           # artificial damping NOT IMPLEMENTED
         smooth_box            = False         # whether to smooth over the edge of the linear box
         CFL                   = 0.5           # Courant stability factor (require <0.5)
-        scale_box_l           = 1.0           # linear box length scale factor in radial direction (left)
-        scale_box_r           = 1.0           # linear box length scale factor in radial direction (right)
-        scale_box_ang_t       = 1.0           # linear box length scale factor in angular direction (top)
-        scale_box_ang_b       = 1.0           # linear box length scale factor in angular direction (bottom)
+        scale_box_left        = 1.0           # linear box length scale factor in radial direction (left)
+        scale_box_right       = 1.0           # linear box length scale factor in radial direction (right)
+        scale_box_ang_top     = 1.0           # linear box length scale factor in angular direction (top)
+        scale_box_ang_bottom  = 1.0           # linear box length scale factor in angular direction (bottom)
         tf_fac                = 1.0           # scale factor for t coordinate where self-similar solution is used
         show_teta_debug_plots = False         # show (t,eta,chi) space developer plots
         box_warp              = True          # interpret y coordinate of linear regime as arc length, or truly vertical? True (default) for former
-        use_box_IC            = False         # use only part of linear regime in box as initial condition for non-linear evolution
+        use_box_IC            = False         # use linear regime in square box as initial condition for non-linear evolution
         use_old_vel           = False         # use old approximated formulas for u pert
+        rot_interp            = False         # expand grid to avoid border effects when interpolatin on rotated grid
+        lin_type              = "global"      # Choose the perturbations to use in the linear regime. Supported options: global, simulation, shearing_sheet
+        nl_wake               = False         # Add non linear correction to wake structure
+        r_cut_inner_fac       = 2             # Cut numerical solution at r=r_planet/r_cut_in_fac
 
         # generate dictionary for model parameters by grabbing all local variables
         self.model_params = locals()
@@ -212,7 +220,7 @@ def configure_from_file(
         print(f"Model configuration read from file: {param_file}")
 
     # generate the model using the configuration specified by the user
-    def run(self, overwrite: bool = False) -> None:
+    def run(self, overwrite: bool = False) -> Optional[float]:
         """
         Generate results for model, requires user to have called either configure or configure_from_file first.
 
@@ -261,7 +269,7 @@ def run(self, overwrite: bool = False) -> None:
 
     # internal method that is called by self.run to generate the results for a specific set of parameters
     # may be called more than once if the user has specified multiple planet masses
-    def _run_wakeflow(self, params: _Parameters) -> None:
+    def _run_wakeflow(self, params: _Parameters) -> Optional[float]:
         """
         Internal use method for generating the planet wake by calling other parts of Wakeflow.
 
@@ -291,7 +299,7 @@ def _run_wakeflow(self, params: _Parameters) -> None:
 
             # extract linear perturbations from file
             lin_perts = _LinearPerts(params)
-            lin_perts._cut_box_annulus_segment()
+            lin_perts._cut_annulus_segment()
 
             # add the linear perturbations onto grid
             grid_lin_perts._add_linear_perturbations(lin_perts, grid_background.rho)
@@ -324,10 +332,11 @@ def _run_wakeflow(self, params: _Parameters) -> None:
             # initialise non-linear perturbations
             nonlin_perts = _NonLinearPerts(params, grid_nonlin_perts)
 
-            # extract initial condition from the linear perturbations
-            nonlin_perts._extract_ICs(lin_perts)
+            # extract initial condition from the linear perturbations using annulus segment
+            nonlin_perts._extract_ICs_ann(lin_perts)
+            # extract initial condition from the linear perturbations using square box
             if params.use_box_IC:
-                nonlin_perts._extract_ICs_ann(lin_perts)
+                nonlin_perts._extract_ICs(lin_perts)
 
             # solve for non-linear perturbations
             nonlin_perts._get_non_linear_perts()


Test patch:
diff --git a/tests/test_model_setup.py b/tests/test_model_setup.py
index 2c28a55..2d28c30 100644
--- a/tests/test_model_setup.py
+++ b/tests/test_model_setup.py
@@ -253,6 +253,23 @@ def test_show_teta_warning(capfd):
     captured = capfd.readouterr()
     assert "WARNING: Choosing show_teta_debug_plots=True may cause the run to fail." in captured.out
 
+def test_mcmc_warning(capfd):
+    
+    # make params dict
+    model = WakeflowModel()
+    model.configure()
+    param_dict = model.model_params
+
+    # use mcmc
+    param_dict["mcmc"] = True
+    params = _Parameters(param_dict)
+
+    # check that sanity checks warns about saving perturbations during mcmcm
+    params._do_sanity_checks()
+    captured = capfd.readouterr()
+    assert "WARNING: Do not save the perturbation fields when using wakeflow inside the mcmc chain" in captured.out
+    assert "WARNING: Do not save the density and velocity fields when using wakeflow inside the mcmc chain" in captured.out
+
 def test_load_config_file():
 
     # make default params dict
diff --git a/tests/test_wakeflow.py b/tests/test_wakeflow.py
index 0dde7f6..9b3ab94 100644
--- a/tests/test_wakeflow.py
+++ b/tests/test_wakeflow.py
@@ -42,6 +42,7 @@ def test_default_user_config():
     assert model.model_params["include_linear"]      == True
     assert model.model_params["save_perturbations"]  == True
     assert model.model_params["save_total"]          == True
+    assert model.model_params["rot_interp"]          == False
     assert model.model_params["write_FITS"]          == False
     assert model.model_params["run_mcfost"]          == False
     assert model.model_params["inclination"]         == -225
@@ -71,6 +72,11 @@ def test_default_dev_config():
     assert model.model_params["show_teta_debug_plots"] == False
     assert model.model_params["box_warp"]              == True
     assert model.model_params["use_box_IC"]            == False
+    assert model.model_params["use_old_vel"]           == False
+    assert model.model_params["rot_interp"]            == False
+    assert model.model_params["mcmc"]                  == False
+    assert model.model_params["nl_wake"]               == False 
+    assert model.model_params["lin_type"]              == "global"
 
 # test that user configuration overrides default
 def test_manual_config():