Einarsson15.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Nov  7 19:15:57 2025

@author: davide
"""


import numpy as np
from scipy import interpolate
from scipy.integrate import odeint
from scipy.optimize import curve_fit


def deltaC_einarsson15(r,Rep,dir='',IC=None,time=None,gammadot=None):
    """returns rotational dinamics from initial condition IC along time at Reynolds Rep for the spheroid with aspect ratio"""
    ###load the coefficients
    b1,b2,b3,b4 = load_betas(dir,req=r,Rep=Rep)
    ###get the prediction
    nnn,ttt = Einarsson_interface(IC=IC,time=time,gammadot=gammadot,r=r,b1=b1,b2=b2,b3=b3,b4=b4)
    ###split and measure C
    return split_orbits_and_get_C(nnn[:,0],nnn[:,1],nnn[:,2],r)

def load_betas(folder,req,Rep):
    """Loads the graphically extrapolated coefficients from Einarsson et al., PoF 2015, Fig. 2"""
    B1data = np.loadtxt(folder+'B1.txt') #Load the data for the beta_1 coefficient as a two-Dim array with format: r,b_1(r)/Rep
    B2data = np.loadtxt(folder+'B2.txt') #Load the data for the beta_2 coefficient: r,b_2(r)  
    B3data = np.loadtxt(folder+'B3.txt') #Load the data for the beta_3 coefficient: r,b_3(r)
    B4data = np.loadtxt(folder+'B4.txt') #Load the data for the beta_4 coefficient: r,b_4(r)
    ###Define the interpolating function and interpolate the four coefficients at the current aspect ratio
    f1 = interpolate.interp1d(B1data[:,0], B1data[:,1])   #linear interpolation function for b_1
    f2 = interpolate.interp1d(B2data[:,0], B2data[:,1])   #linear interpolation function for b_2
    f3 = interpolate.interp1d(B3data[:,0], B3data[:,1])   #linear interpolation function for b_3
    f4 = interpolate.interp1d(B4data[:,0], B4data[:,1])   #linear interpolation function for b_4
    ###calculate the coefficients at particle Reynolds number Rep
    betas = np.array([f1(req),f2(req),f3(req),f4(req)])* Rep
    return betas

def rotation_Einarsson(p,t,beta,gammadot,b1,b2,b3,b4):
    """Calculates the rotation rate according to Einarsson et al., PoF, 2015: eq. 39"""
    Omega =  np.array([[0, 0.5, 0],[-0.5, 0, 0],[ 0, 0, 0]]) * gammadot     #Rate of rotation tensor
    E = np.array([[0, 0.5, 0],[0.5, 0, 0],[ 0, 0, 0]]) * gammadot           #Rate of strain tensor
    pEp=np.dot(p,np.dot(E,p))                                                        #Usefull term
    pOEp=np.dot(p,np.dot(Omega,np.dot(E,p)))                                         #Usefull term
    pEEp=np.dot(p,np.dot(E,np.dot(E,p))) 
    b1/=gammadot;b2/=gammadot;b3/=gammadot;b4/=gammadot                                            #Usefull term
    return np.dot(Omega+beta*E,p)-beta*pEp*p+b1*pEp*(np.dot(E,p)-p*pEp)+b2*pEp*np.dot(Omega,p)+b3*(np.dot(Omega,np.dot(E,p))-pOEp*p)+b4*(np.dot(E,np.dot(E,p))-pEEp*p)

def Einarsson_interface(IC,time,gammadot,r,b1,b2,b3,b4):
    """Manages the integration of the Einarsson rotation rate from a given initial_value"""
    beta = (r**2-1)/(r**2+1)                                                              #Calculate the geometrical parameter beta
    simtime = time #time-range of integration
    pinit=np.array([IC[0],IC[1],IC[2]])         ###initial conditions q0
    ptheory = odeint(rotation_Einarsson,pinit,simtime,args=(beta,gammadot,b1,b2,b3,b4)) # solve the problem using odeint and calling the previously defined function rotation_Einarsson
    return ptheory,simtime #return the direction vector and the integration time 



def jeffery_model(phi, C, r):
    return C * r / np.sqrt(np.cos(phi)**2 + r**2 * np.sin(phi)**2)

def fit_C(theta,phi,r):
    def model(phi, C):
        return jeffery_model(phi, C, r)
    # Fit
    try:
        C_fit, cov = curve_fit(model, phi, np.tan(theta), bounds=[0, np.inf], p0=[1.0])
        C_std = np.sqrt(np.diag(cov))[0] if cov.size else np.nan
        return C_fit[0], C_std
    except Exception:
            return np.nan, np.nan

def split_orbits_and_get_C(n1,n2,n3,r):
    """Method to split orbits in semi-periods"""
    # Identify where sign changes
    if r > 1.0:
        sign_change = np.diff(np.sign(n1)) != 0
    else:
        sign_change = np.diff(np.sign(n2)) != 0
    orbit_labels = np.zeros_like(n1, dtype=int)
    orbit_labels[1:] = np.cumsum(sign_change)
    num_orbits = orbit_labels.max() + 1
    C = []
    for i in range(num_orbits):
        orbit_inds = np.where(orbit_labels[:]==i)
        # phi = np.arctan(n1[orbit_inds]/n2[orbit_inds])
        phi = np.arctan2(n1[orbit_inds],n2[orbit_inds])
        theta = np.arccos(n3[orbit_inds])
        # tan_theta = np.tan(theta)
        C.append(fit_C(theta,phi,r))
    return np.array(C).T