Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Improve scipy.integrate.quad_vec performance for fitting integral equation, workers keyword non-functional

I am using quad_vec in a function that contains an integral not solvable with analytic methods. I need to fit this equation to data points.

However, even the evaluation with fixed values takes multiple seconds, the total fit time grows to 7 min.

Is there a way to accelerate the computation? I know quad_vec has the workers keyword, but when I try to use that, the computation does not finish at all. I am currently working in a Jupyter notebook if that has any significance.

This is the function definition, I already tried to use numba here with little success.

from scipy.special import j0, j1
from scipy.integrate import quad_vec
import numpy as np
import numba as nb

h = 0.00005
G = 1000
E = 3*G
R= 0.0002
upsilon = 0.06
gamma = 0.072

@nb.jit(nb.float64(nb.float64, nb.float64, nb.float64, nb.float64, nb.float64),cache=True)
def QSzz_1(s,h,z, upsilon, E):
    # exponential form of the QSzz^-1 function to prevent float overflow,
    # also assumes nu==0.5
    numerator = (1+2*s**2*h**2)*np.exp(-2*s*z) + 0.5*(1+np.exp(-4*s*z))
    denominator = 0.5*(1-np.exp(-4*s*z)) - 2*s*h*np.exp(-2*s*z)
    return (3/(2*s*E)) / (numerator/denominator + (3*upsilon/(2*E))*s)

def integrand(s, r, R, upsilon, E, h):
    return s*(R*j0(s*R) - 2*j1(s*R)/s) * QSzz_1(s,h,h, upsilon, E) * j0(s*r)

def style_exact(r, gamma, R, upsilon, E, h):               
    int_out = quad_vec(integrand, 0, np.inf, args=(r,R,upsilon,E,h))
    return gamma * int_out[0]

# calculate fixed ~10s
x_ax = np.linspace(0, 0.0004,101, endpoint=True, dtype=np.float64)
zeta = style_exact(x_ax, gamma, 0.0002, upsilon, E, h)

# fit to dataset (wetting ridge) ~7 min
popt_hemi, pcov_hemi = curve_fit(lambda x, upsilon, E, h: style_exact(x, gamma, R, upsilon, E, h) ,points_x, points_y, p0=(upsilon, E, h), bounds=([0,0,0],[np.inf,np.inf,np.inf]))

Edit: here are some exemplar values:

points_x =[0.00040030286, 0.00040155788, 0.00040281289, 0.00040406791000000003, 0.00040532292, 0.00040657794, 0.00040783296, 0.00040908797, 0.00041034299, 0.00041159801, 0.00041285302, 0.00041410804, 0.00041536305, 0.00041661807, 0.00041787309, 0.0004191281, 0.00042038312, 0.00042163814, 0.00042289315000000003, 0.00042414817000000003, 0.00042540318, 0.0004266582, 0.00042791322, 0.00042916823, 0.00043042325, 0.00043167827, 0.00043293328, 0.0004341883, 0.00043544332, 0.00043669833, 0.00043795335, 0.00043920836, 0.00044046338, 0.0004417184, 0.00044297341000000003, 0.00044422843000000003, 0.00044548345, 0.00044673846, 0.00044799348, 0.00044924849, 0.00045050351, 0.00045175852999999996, 0.00045301354000000006, 0.00045426856, 0.00045552357999999995, 0.00045677859000000005, 0.00045803361, 0.00045928863000000006, 0.00046054364000000004, 0.00046179866, 0.00046305367, 0.00046430869000000004, 0.00046556371, 0.00046681871999999997, 0.00046807374000000003, 0.00046932876, 0.00047058376999999996, 0.00047183879, 0.0004730938, 0.00047434881999999995, 0.00047560384, 0.00047685885, 0.00047811387000000006, 0.00047936889, 0.0004806239, 0.00048187892000000005, 0.00048313393000000004, 0.00048438895, 0.00048564397000000004, 0.00048689898000000003, 0.000488154, 0.00048940902, 0.00049066403, 0.00049191905, 0.00049317407, 0.00049442908, 0.0004956841, 0.0004969391100000001, 0.00049819413, 0.0004994491500000001, 0.00050070416, 0.00050195918, 0.0005032142, 0.00050446921, 0.00050572423, 0.00050697924, 0.00050823426, 0.00050948928, 0.00051074429, 0.00051199931, 0.00051325433, 0.00051450934, 0.00051576436, 0.00051701937, 0.0005182743900000001, 0.00051952941, 0.00052078442, 0.00052203944, 0.00052329446, 0.00052454947, 0.00052580449, 0.00052705951, 0.00052831452, 0.00052956954, 0.00053082455, 0.00053207957, 0.00053333459, 0.0005345896, 0.00053584462, 0.00053709964, 0.00053835465, 0.00053960967, 0.00054086468, 0.0005421197, 0.0005433747200000001, 0.00054462973, 0.00054588475, 0.00054713977, 0.00054839478, 0.0005496498, 0.00055090482, 0.00055215983, 0.00055341485, 0.00055466986, 0.00055592488, 0.0005571799, 0.00055843491, 0.00055968993, 0.00056094495, 0.0005621999600000001, 0.00056345498, 0.00056470999, 0.00056596501, 0.00056722003, 0.00056847504, 0.00056973006, 0.00057098508, 0.00057224009, 0.00057349511, 0.00057475012, 0.00057600514, 0.00057726016, 0.00057851517, 0.00057977019, 0.00058102521, 0.00058228022, 0.0005835352400000001, 0.00058479026, 0.00058604527, 0.00058730029, 0.0005885553, 0.00058981032, 0.00059106534, 0.00059232035, 0.00059357537, 0.00059483039, 0.0005960854, 0.00059734042, 0.00059859543, 0.00059985045, 0.00060110547, 0.0006023604800000001, 0.0006036155, 0.00060487052, 0.00060612553, 0.00060738055, 0.00060863556, 0.00060989058, 0.0006111456, 0.00061240061, 0.00061365563, 0.00061491065, 0.00061616566, 0.00061742068, 0.0006186757, 0.00061993071, 0.00062118573, 0.00062244074, 0.00062369576, 0.00062495078, 0.00062620579, 0.0006274608100000001, 0.00062871583, 0.00062997084, 0.00063122586, 0.00063248087, 0.00063373589, 0.00063499091, 0.00063624592, 0.00063750094, 0.00063875596, 0.00064001097, 0.00064126599, 0.000642521, 0.00064377602, 0.00064503104, 0.0006462860500000001, 0.00064754107, 0.0006487960900000001, 0.0006500511, 0.00065130612, 0.00065256114, 0.00065381615, 0.00065507117, 0.00065632618, 0.0006575812, 0.00065883622, 0.00066009123, 0.00066134625, 0.00066260127, 0.00066385628, 0.0006651113, 0.00066636631, 0.0006676213300000001, 0.00066887635, 0.00067013136, 0.00067138638, 0.0006726414, 0.00067389641, 0.00067515143, 0.00067640645, 0.00067766146, 0.00067891648, 0.00068017149, 0.00068142651, 0.00068268153, 0.00068393654, 0.00068519156, 0.00068644658, 0.00068770159, 0.00068895661, 0.00069021162, 0.00069146664, 0.0006927216600000001, 0.00069397667, 0.00069523169, 0.00069648671, 0.00069774172, 0.00069899674, 0.00070025175, 0.00070150677, 0.00070276179, 0.0007040168, 0.00070527182, 0.00070652684, 0.00070778185, 0.00070903687, 0.00071029189, 0.0007115469000000001, 0.00071280192, 0.00071405693, 0.00071531195, 0.00071656697, 0.00071782198, 0.000719077, 0.00072033202, 0.00072158703, 0.00072284205, 0.00072409706, 0.00072535208, 0.0007266071, 0.00072786211, 0.00072911713, 0.00073037215, 0.00073162716, 0.0007328821800000001, 0.00073413719, 0.00073539221, 0.00073664723, 0.00073790224, 0.00073915726, 0.00074041228, 0.00074166729, 0.00074292231, 0.00074417733, 0.00074543234, 0.00074668736, 0.00074794237, 0.00074919739, 0.00075045241, 0.0007517074200000001, 0.00075296244, 0.00075421746, 0.00075547247, 0.00075672749, 0.0007579825, 0.00075923752, 0.00076049254, 0.00076174755, 0.00076300257, 0.00076425759, 0.0007655126, 0.00076676762, 0.00076802264, 0.00076927765, 0.00077053267, 0.00077178768, 0.0007730427, 0.00077429772, 0.00077555273, 0.0007768077500000001, 0.00077806277, 0.00077931778, 0.0007805728, 0.00078182781, 0.00078308283, 0.00078433785, 0.00078559286, 0.00078684788, 0.0007881029, 0.00078935791, 0.00079061293, 0.00079186794, 0.00079312296, 0.00079437798, 0.0007956329900000001, 0.00079688801, 0.0007981430300000001, 0.00079939804, 0.00080065306, 0.00080190808, 0.00080316309, 0.00080441811, 0.00080567312, 0.00080692814, 0.00080818316, 0.00080943817, 0.00081069319, 0.00081194821, 0.00081320322, 0.00081445824, 0.00081571325, 0.0008169682700000001, 0.00081822329, 0.0008194783, 0.00082073332, 0.00082198834, 0.00082324335, 0.00082449837, 0.00082575338, 0.0008270084, 0.00082826342, 0.00082951843, 0.00083077345, 0.00083202847, 0.00083328348, 0.0008345385, 0.00083579352, 0.00083704853, 0.00083830355, 0.00083955856, 0.00084081358, 0.0008420686000000001, 0.00084332361, 0.00084457863, 0.00084583365, 0.00084708866, 0.00084834368, 0.00084959869, 0.00085085371, 0.00085210873, 0.00085336374, 0.00085461876, 0.00085587378, 0.00085712879, 0.00085838381, 0.00085963882, 0.0008608938400000001, 0.00086214886, 0.00086340387, 0.00086465889, 0.00086591391, 0.00086716892, 0.00086842394, 0.00086967896, 0.00087093397, 0.00087218899, 0.000873444, 0.00087469902, 0.00087595404, 0.00087720905, 0.00087846407, 0.00087971909, 0.0008809741, 0.0008822291200000001, 0.00088348413, 0.00088473915, 0.00088599417, 0.00088724918, 0.0008885042, 0.00088975922, 0.00089101423, 0.00089226925, 0.00089352427, 0.00089477928, 0.0008960343, 0.00089728931, 0.00089854433, 0.00089979935, 0.0009010543600000001, 0.00090230938, 0.0009035644, 0.00090481941, 0.00090607443, 0.00090732944, 0.00090858446, 0.00090983948, 0.00091109449, 0.00091234951, 0.00091360453, 0.00091485954, 0.00091611456, 0.00091736957, 0.00091862459, 0.00091987961, 0.00092113462, 0.00092238964, 0.00092364466, 0.00092489967, 0.0009261546900000001, 0.00092740971, 0.00092866472, 0.00092991974, 0.00093117475, 0.00093242977, 0.00093368479, 0.0009349398, 0.00093619482, 0.00093744984, 0.00093870485, 0.0009399598700000001, 0.00094121488, 0.0009424699, 0.0009437249200000001, 0.00094497993, 0.00094623495, 0.0009474899700000001, 0.0009487449799999999, 0.00095, 0.0009512550100000001, 0.0009525100299999999, 0.00095376505, 0.0009550200600000001, 0.0009562750799999999, 0.0009575301, 0.0009587851100000001, 0.0009600401299999999, 0.00096129515, 0.00096255017, 0.00096380517, 0.0009650601700000001, 0.00096631517, 0.00096757027, 0.0009688252700000001, 0.00097008027, 0.0009713352699999999, 0.0009725902700000001, 0.00097384527, 0.0009751003699999999, 0.0009763553700000001, 0.00097761037, 0.00097886537, 0.00098012037, 0.0009813753699999999, 0.0009826304699999998, 0.0009838854700000002, 0.00098514047, 0.00098639547, 0.00098765047, 0.0009889054699999998, 0.0009901605699999998, 0.0009914155700000002, 0.00099267057, 0.00099392557, 0.00099518057, 0.0009964355699999998, 0.0009976905700000002, 0.0009989456700000001, 0.00100020067, 0.00100145567, 0.00100271067, 0.0010039656699999998, 0.0010052206700000002, 0.0010064757700000001, 0.00100773077, 0.00100898577, 0.0010102407699999999, 0.0010114957699999998, 0.0010127507700000002, 0.00101400587, 0.00101526087, 0.00101651587, 0.0010177708699999999, 0.0010190258700000002, 0.0010202808700000001, 0.00102153597, 0.00102279097, 0.00102404597, 0.0010253009699999998, 0.0010265559700000002, 0.0010278109700000001, 0.00102906607, 0.00103032107, 0.00103157607, 0.0010328310699999998, 0.0010340860700000002, 0.00103534107, 0.00103659617, 0.00103785117, 0.00103910617, 0.0010403611699999998, 0.0010416161700000002, 0.00104287117, 0.00104412627, 0.00104538127, 0.0010466362699999999, 0.0010478912699999998, 0.0010491462700000002, 0.00105040127, 0.00105165627, 0.00105291137, 0.0010541663699999999, 0.0010554213700000002, 0.0010566763700000001, 0.00105793137, 0.00105918637, 0.00106044147, 0.0010616964699999999, 0.0010629514700000002, 0.0010642064700000001, 0.00106546147, 0.00106671647, 0.00106797157, 0.0010692265699999998, 0.0010704815700000002, 0.00107173657, 0.00107299157, 0.00107424657, 0.00107550167, 0.0010767566699999998, 0.0010780116700000002, 0.00107926667, 0.00108052167, 0.00108177667, 0.0010830317699999999, 0.0010842867699999998, 0.0010855417700000002, 0.00108679677, 0.00108805177, 0.00108930677, 0.0010905618699999999, 0.0010918168700000002, 0.0010930718700000001, 0.00109432687, 0.00109558187, 0.00109683687, 0.0010980919699999999, 0.0010993469700000002, 0.0011006019700000001, 0.00110185697, 0.00110311197, 0.0011043669699999999, 0.0011056219699999998, 0.0011068770700000002, 0.00110813207, 0.00110938707, 0.00111064207, 0.0011118970699999999, 0.0011131520700000002, 0.0011144071700000002, 0.00111566217, 0.00111691717, 0.00111817217, 0.0011194271699999998, 0.0011206821700000002, 0.0011219372700000002, 0.00112319227, 0.00112444727, 0.00112570227, 0.0011269572699999998, 0.0011282122700000002, 0.0011294673700000001, 0.00113072237, 0.00113197737, 0.00113323237, 0.0011344873699999998, 0.0011357423700000002, 0.0011369974700000001, 0.00113825247, 0.00113950747, 0.0011407624699999999, 0.0011420174699999998, 0.0011432724700000002, 0.0011445275700000001, 0.00114578257, 0.00114703757, 0.0011482925699999999, 0.0011495475700000002, 0.0011508025700000001, 0.00115205767, 0.00115331267, 0.00115456767, 0.0011558226699999999, 0.0011570776700000002, 0.0011583326700000001, 0.00115958777, 0.00116084277, 0.00116209777, 0.0011633527699999998, 0.0011646077700000002, 0.00116586277, 0.00116711777, 0.00116837287, 0.00116962787, 0.0011708828699999998, 0.0011721378700000002, 0.00117339287, 0.00117464787, 0.00117590297, 0.0011771579699999999, 0.0011784129699999998, 0.0011796679700000002, 0.00118092297, 0.00118217797, 0.00118343307, 0.0011846880699999999, 0.0011859430700000002, 0.0011871980700000001, 0.00118845307, 0.00118970807, 0.00119096317, 0.0011922181699999999, 0.0011934731700000002, 0.0011947281700000001, 0.00119598317, 0.00119723817, 0.00119849327, 0.0011997482699999998, 0.0012010032700000002, 0.00120225827, 0.00120351327, 0.00120476827, 0.00120602337, 0.0012072783699999998, 0.0012085333700000002, 0.00120978837, 0.00121104337, 0.00121229837, 0.0012135534699999999, 0.0012148084699999998, 0.0012160634700000002, 0.00121731847, 0.00121857347, 0.00121982847, 0.0012210834699999998, 0.0012223385700000002, 0.0012235935700000001, 0.00122484857, 0.00122610357, 0.00122735857, 0.0012286135699999998, 0.0012298686700000002, 0.0012311236700000001, 0.00123237867, 0.00123363367, 0.0012348886699999999, 0.0012361436699999998]


points_y = [-2.4929826e-07, -2.3248189e-07, -4.4305314e-07, -1.0689171e-06, -7.0144722e-07, -1.3773717e-06, -9.3672285e-07, -1.6876499e-06, -9.8346007e-07, -1.7992562e-06, -1.0198111e-06, -1.721233e-06, -8.9082583e-07, -1.1925362e-06, -8.3776501e-07, -6.9998957e-07, -7.1134901e-07, -4.5476849e-07, -6.4449894e-07, -3.8765887e-07, -6.3044764e-07, -7.4224008e-07, -7.6114851e-07, -1.0377502e-06, -1.3589471e-06, -1.3342596e-06, -1.3712255e-06, -1.3510569e-06, -1.2278933e-06, -8.2319036e-07, -1.4040568e-06, -6.4183121e-07, -1.1649824e-06, -7.3197454e-07, -1.0537769e-06, -8.3223932e-07, -1.1644648e-06, -1.2177416e-06, -1.4045247e-06, -1.6934001e-06, -1.6157397e-06, -1.8595331e-06, -1.7097882e-06, -1.8031869e-06, -1.5406345e-06, -1.5851084e-06, -1.5695719e-06, -1.4990693e-06, -1.8087735e-06, -1.7151045e-06, -1.8353234e-06, -1.71844e-06, -1.7904118e-06, -1.5297879e-06, -1.6064767e-06, -1.4520618e-06, -1.1090131e-06, -1.2475477e-06, -7.4591269e-07, -1.0619496e-06, -7.5699762e-07, -1.3883064e-06, -1.3300594e-06, -1.9713711e-06, -2.0613271e-06, -2.5116161e-06, -2.4466345e-06, -2.5386926e-06, -2.2368298e-06, -2.2934508e-06, -1.8951084e-06, -1.8117756e-06, -1.6680112e-06, -1.8274169e-06, -1.7569355e-06, -2.1081536e-06, -2.1241154e-06, -2.2742958e-06, -2.4032149e-06, -2.2596226e-06, -2.1889918e-06, -1.9359605e-06, -1.8878718e-06, -1.6144539e-06, -1.6485844e-06, -1.2316506e-06, -1.6932815e-06, -8.1348768e-07, -1.310099e-06, -4.3574574e-07, -1.0726973e-06, -6.6005902e-07, -1.2151841e-06, -9.1100721e-07, -1.4911344e-06, -1.3152027e-06, -1.3695714e-06, -1.3930563e-06, -1.3452594e-06, -1.3228626e-06, -1.3714694e-06, -1.2480971e-06, -1.4622823e-06, -1.5687181e-06, -1.7872703e-06, -1.7135845e-06, -2.0209804e-06, -1.3665688e-06, -1.7074398e-06, -1.1511678e-06, -1.1604734e-06, -1.0173458e-06, -1.0660268e-06, -1.0424449e-06, -1.1101976e-06, -1.0030326e-06, -1.0879421e-06, -8.2978143e-07, -9.3823628e-07, -7.2342249e-07, -9.8929055e-07, -1.0764783e-06, -1.3105722e-06, -1.3954326e-06, -1.5047949e-06, -1.4339143e-06, -1.3061363e-06, -1.3200332e-06, -1.381963e-06, -1.3490984e-06, -1.3526509e-06, -1.463083e-06, -1.2588114e-06, -1.4445926e-06, -1.1240129e-06, -1.3659935e-06, -1.3323392e-06, -1.3695779e-06, -1.7108472e-06, -1.7111548e-06, -2.0250494e-06, -2.1803196e-06, -2.2433208e-06, -2.4435685e-06, -1.9341618e-06, -2.3866277e-06, -1.8497934e-06, -1.8903583e-06, -1.4422203e-06, -1.7661343e-06, -1.5059728e-06, -1.5770287e-06, -1.8108199e-06, -2.0170832e-06, -1.8260586e-06, -2.1429269e-06, -2.0532939e-06, -2.1373399e-06, -2.342127e-06, -2.3871293e-06, -2.5980083e-06, -2.4293864e-06, -2.3568741e-06, -2.0801477e-06, -1.8587702e-06, -1.7074138e-06, -1.5791169e-06, -1.6891695e-06, -1.7635139e-06, -1.9566623e-06, -1.8455385e-06, -2.1080438e-06, -2.0320153e-06, -2.1665641e-06, -2.1571212e-06, -2.3643005e-06, -2.074037e-06, -2.0893195e-06, -1.9232214e-06, -1.7025658e-06, -1.6232691e-06, -1.6510243e-06, -1.7197265e-06, -1.8580166e-06, -1.9258182e-06, -2.0062691e-06, -2.0157544e-06, -2.0394525e-06, -2.0826713e-06, -1.9067459e-06, -2.0218438e-06, -1.9964327e-06, -2.1734356e-06, -2.1242189e-06, -2.4424379e-06, -2.4437198e-06, -2.6022861e-06, -2.4502697e-06, -2.6343237e-06, -2.2225432e-06, -2.3110892e-06, -2.1664638e-06, -2.1287713e-06, -2.011825e-06, -2.2808875e-06, -2.158988e-06, -2.5522458e-06, -2.556647e-06, -2.8299596e-06, -2.9620166e-06, -2.6908558e-06, -3.0163631e-06, -2.6530144e-06, -2.5642676e-06, -2.2324086e-06, -2.0825715e-06, -1.7085644e-06, -1.4025919e-06, -1.4042667e-06, -1.397307e-06, -1.4471031e-06, -1.4352464e-06, -1.6847902e-06, -1.4372545e-06, -1.6405646e-06, -1.5025385e-06, -1.58785e-06, -1.5018164e-06, -1.546755e-06, -1.5307927e-06, -1.5450872e-06, -1.762507e-06, -1.9245396e-06, -2.1342847e-06, -2.083201e-06, -2.1824533e-06, -2.2264199e-06, -1.9521925e-06, -2.1104425e-06, -2.35205e-06, -2.1372429e-06, -2.3874246e-06, -2.3111549e-06, -2.3476044e-06, -1.9828263e-06, -2.1105666e-06, -1.77767e-06, -1.8420129e-06, -1.90373e-06, -1.930438e-06, -2.0727705e-06, -2.1793671e-06, -2.4205829e-06, -2.1275047e-06, -2.4740434e-06, -2.0603233e-06, -2.2409819e-06, -1.7541814e-06, -2.0279909e-06, -1.730486e-06, -1.9476207e-06, -1.7534857e-06, -1.8505329e-06, -2.0095086e-06, -1.618978e-06, -1.8867553e-06, -1.9088163e-06, -1.886491e-06, -1.7468138e-06, -1.8476389e-06, -1.7557932e-06, -1.4058452e-06, -1.6067978e-06, -1.3156005e-06, -1.3659535e-06, -1.0961384e-06, -1.0153987e-06, -9.4432646e-07, -6.6454642e-07, -9.2586387e-07, -1.0025458e-06, -1.0698426e-06, -1.2805659e-06, -1.3957816e-06, -1.504749e-06, -1.3274602e-06, -1.4140738e-06, -1.3504825e-06, -1.3899331e-06, -1.3970904e-06, -1.4744283e-06, -1.4185692e-06, -1.7050143e-06, -1.5382651e-06, -1.5599202e-06, -1.5529446e-06, -1.506719e-06, -1.4330019e-06, -1.240627e-06, -1.2835575e-06, -1.1023492e-06, -1.1632735e-06, -1.1683113e-06, -1.2732747e-06, -1.219676e-06, -1.2890147e-06, -1.1440703e-06, -9.1523203e-07, -8.2035542e-07, -8.7226368e-07, -7.3633722e-07, -9.884313e-07, -8.5961273e-07, -1.2392311e-06, -1.0843573e-06, -1.0707268e-06, -9.571558e-07, -1.0067944e-06, -6.4553431e-07, -4.0506156e-07, -3.3043043e-07, -1.7361598e-07, -1.3118263e-07, -2.9468891e-07, -4.7080768e-07, -6.4225818e-07, -7.5475209e-07, -8.5102358e-07, -6.0803728e-07, -8.1677753e-07, -5.9744241e-07, -7.8274568e-07, -5.3968306e-07, -8.350585e-07, -5.4845851e-07, -5.8427222e-07, -5.1520419e-07, -4.6822083e-07, -6.0910398e-07, -4.4298342e-07, -5.6257054e-07, -2.7562129e-07, -2.5181401e-07, 3.8053095e-08, 2.4159147e-07, 3.6882074e-07, 5.1241897e-07, 4.8644598e-07, 7.2692073e-07, 4.7022181e-07, 7.0384493e-07, 6.8289479e-07, 6.4066943e-07, 8.5657662e-07, 5.8406311e-07, 6.7344028e-07, 4.1435118e-07, 2.7649325e-07, 2.3123522e-07, -1.9399705e-08, 2.6291987e-07, -3.6143527e-08, 4.1732021e-07, 3.3391364e-07, 6.4314122e-07, 7.7139665e-07, 1.1209136e-06, 1.4367421e-06, 1.6319081e-06, 1.7711259e-06, 1.8566403e-06, 1.7371454e-06, 1.4824876e-06, 1.6134811e-06, 1.0707754e-06, 1.3415844e-06, 1.1356512e-06, 1.4106389e-06, 1.4104486e-06, 1.7408528e-06, 2.0744193e-06, 2.079919e-06, 2.1838213e-06, 2.3656145e-06, 2.1909773e-06, 2.3504607e-06, 2.2917643e-06, 2.4505978e-06, 2.0934847e-06, 2.583584e-06, 2.2871518e-06, 2.5116042e-06, 2.6234818e-06, 2.8420594e-06, 3.0011699e-06, 3.3721137e-06, 3.3177881e-06, 3.7014297e-06, 3.4988464e-06, 3.6346743e-06, 3.6031151e-06, 3.6434367e-06, 3.3825082e-06, 3.6445565e-06, 3.2970635e-06, 3.6138927e-06, 3.3753039e-06, 3.7447733e-06, 3.5673385e-06, 3.6078831e-06, 3.5609168e-06, 3.6213054e-06, 3.5038571e-06, 3.7264648e-06, 3.9751613e-06, 3.8206903e-06, 4.1254495e-06, 3.9272576e-06, 4.2386514e-06, 3.815278e-06, 4.2691643e-06, 4.0643683e-06, 4.330484e-06, 4.29042e-06, 4.6035887e-06, 4.4565016e-06, 4.583597e-06, 4.7192276e-06, 4.7442267e-06, 4.734727e-06, 5.0407053e-06, 5.3132589e-06, 5.3419609e-06, 5.7940368e-06, 6.014359e-06, 6.0453411e-06, 6.0996584e-06, 6.064599e-06, 6.1232403e-06, 5.8926808e-06, 6.0748121e-06, 5.9732831e-06, 6.0281785e-06, 5.9558067e-06, 5.8235522e-06, 5.6378228e-06, 5.4438118e-06, 5.3658419e-06, 5.3454619e-06, 5.205238e-06, 5.4038907e-06, 5.0070169e-06, 5.0996156e-06, 4.5688268e-06, 4.5768204e-06, 4.3706204e-06, 4.378131e-06, 4.3035565e-06, 4.4136234e-06, 4.4586055e-06, 4.2999055e-06, 4.2367521e-06, 4.1092479e-06, 3.6691199e-06, 3.7132548e-06, 3.3891334e-06, 3.4132172e-06, 3.3112791e-06, 3.4194779e-06, 3.3548478e-06, 3.4746562e-06, 3.0714297e-06, 3.631046e-06, 2.9155762e-06, 3.3648723e-06, 3.0564361e-06, 3.1977623e-06, 2.9422311e-06, 2.8664619e-06, 2.9553471e-06, 2.6331467e-06, 2.7458985e-06, 2.4857213e-06, 2.5358048e-06, 2.0853043e-06, 2.2717608e-06, 1.7708539e-06, 2.1185441e-06, 1.8057521e-06, 2.1431184e-06, 1.8050008e-06, 2.2162456e-06, 1.8085417e-06, 2.0822527e-06, 1.6735792e-06, 1.9627324e-06, 1.5854033e-06, 1.7829235e-06, 1.7266717e-06, 1.9015957e-06, 2.1904481e-06, 2.0235789e-06, 2.319506e-06, 2.0939101e-06, 1.9725124e-06, 1.8089637e-06, 1.6690528e-06, 1.5539039e-06, 1.5197157e-06, 1.6846562e-06, 1.5117772e-06, 1.6974785e-06, 1.6371901e-06, 1.62875e-06, 1.3414928e-06, 1.5716781e-06, 1.1625945e-06, 1.4214429e-06, 9.0279233e-07, 1.1886867e-06, 1.0201856e-06, 1.1328523e-06, 1.1236831e-06, 1.2940038e-06, 1.5421363e-06, 1.4063536e-06, 1.8374101e-06, 1.4969428e-06, 1.8342753e-06, 1.3619477e-06, 1.6680712e-06, 1.1305091e-06, 1.3914424e-06, 1.1421031e-06, 1.0667439e-06, 9.2969523e-07, 1.0559697e-06, 9.1449135e-07, 1.0647219e-06, 1.0087653e-06, 1.4041236e-06, 1.0653469e-06, 1.5728387e-06, 1.1900372e-06, 1.396327e-06, 9.5831428e-07, 9.8273849e-07, 6.9228567e-07, 8.1312667e-07, 7.1831973e-07, 8.7380434e-07, 1.1589902e-06, 1.1559309e-06, 1.3702947e-06, 1.2662056e-06, 1.5425231e-06, 1.3134162e-06, 1.3669259e-06, 1.3616054e-06, 1.1954299e-06, 1.2969087e-06, 1.2044857e-06, 1.2129433e-06, 1.066829e-06, 1.2888265e-06, 9.1080301e-07, 9.4850302e-07, 5.628118e-07, 7.3976011e-07, 2.4812591e-07, 4.8428843e-07, 2.8183855e-07, 5.7697958e-07, 4.5050618e-07, 8.9398675e-07, 6.1835844e-07, 1.0104357e-06, 8.0297383e-07, 1.048894e-06, 8.5358493e-07, 1.0796489e-06, 8.1287327e-07, 9.8257456e-07, 5.680762e-07, 8.0814245e-07, 5.9584009e-07, 6.3797485e-07, 6.944886e-07, 9.4480563e-07, 9.518683e-07, 1.2123149e-06, 1.5256706e-06, 1.5178348e-06, 1.5515784e-06, 1.7937264e-06, 1.240253e-06, 1.3527793e-06, 1.0316758e-06, 9.3243026e-07, 9.332284e-07, 7.1470557e-07, 8.428729e-07, 6.2857809e-07, 5.4179424e-07, 4.6470621e-07, 3.6276441e-07, 2.814264e-07, 3.295977e-07, 4.0521404e-07, 3.5343026e-07, 3.6958044e-07, 5.0506593e-07, 3.1834025e-07, 3.2582213e-07, 4.3522668e-07, 2.6025184e-07, 3.8900153e-07, 1.0961338e-07, 2.5467694e-07, 1.166893e-07, 5.6772224e-08, 9.1470554e-08, 2.0167496e-07, 3.7911797e-08, 2.6796461e-07, 2.0933361e-07, 2.1677593e-07, 1.5076751e-07, 2.154547e-07, 9.4922825e-08, -1.5619153e-08, 5.6953286e-08, 1.492038e-08, -1.2234541e-07, -7.3945498e-08, -1.4066427e-07, -1.5021338e-07, -8.5765791e-08, -2.5937592e-07, -8.5784093e-08, -3.7865655e-07, -3.3939569e-07, -5.3969692e-07, -6.6329776e-07, -6.7695552e-07, -7.9978318e-07, -6.5715392e-07, -5.8634763e-07, -3.4631253e-07, -2.6236251e-07, -9.1816048e-09, -6.8072671e-08, 4.6884891e-08, -2.1581414e-07, -1.6978596e-07, -3.1446355e-07, -3.5427569e-07, -2.7410849e-07, -2.8441695e-07, -2.4303658e-07, -6.8944944e-08, -1.8188616e-07, 5.0232102e-08, -5.0489499e-08, -3.4827404e-08, -2.0914572e-07, -2.141703e-07]

like image 439
Raphael Avatar asked Oct 22 '25 20:10

Raphael


1 Answers

There are certainly some code optimizations in the comments that can speed things up, so I'll complement that by focusing on the integration. If you can use a GPU, you can get two to three orders of magnitude speedup there.

One of the simplest things that you can do to improve the speed without sacrificing too much accuracy is to provide limits of integration that capture the important part of the integral. I plotted the integrand to choose 1e7 as the upper limit, but you can probably develop a heuristic and write something to automatically choose a reasonable value.

args= (x_ax, R, upsilon, E, h)

# infinite upper limit of integration
%timeit -r 1 -n 1 quad_vec(integrand, 0, np.inf, args=args)
# 9.41 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
res = quad_vec(integrand, 0, np.inf, args=args)

# large, finite upper limit of integration
%timeit quad_vec(integrand, 0, 1e7, args=args)
# 595 ms ± 169 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
res2 = quad_vec(integrand, 0, 1e7 args=args)

# Check error between the two
# (note that `quad_vec` may not be super accurate, anyway;
#  consider checking against `mpmath.quad`)
from scipy import stats
res = quad_vec(integrand, 0, 1e7, args=args)
res0 = quad_vec(integrand, 0, np.inf, args=args)
stats.gmean(abs((res[0] - res0[0])/res0[0]))  # gmean of relative error
# 3.643331971383965e-05

You can get much bigger gains without too much trouble by rolling your own integration scheme in CuPy (or find an integration library that makes use of the GPU).

For instance, there is an implementation of the tanh-sinh quadrature rule in SciPy that is relatively simple to adapt to something that uses the GPU, especially if you're willing to go with a fixed step size. tanh-sinh is probably not the best rule for this oscillatory integrand; I'm just using it because I'm familiar with it.

import cupy as np  # still np to avoid changing the existing code
from cupyx.scipy.special import j0, j1

# copy-pasted from
# https://github.com/scipy/scipy/blob/v1.12.0/scipy/integrate/_tanhsinh.py
# then removed comments to reduce length; other optimizations possible
# There are plans to make this Array-API compatible and provide a public 
# interface

_N_BASE_STEPS = 8

def _get_base_step(dtype=np.float64):
    fmin = 4*np.finfo(dtype).tiny
    tmax = np.arcsinh(np.log(2/fmin - 1) / np.pi)
    h0 = tmax / _N_BASE_STEPS
    return h0.astype(dtype)

def _compute_pair(k, h0):
    h = h0 / 2**k
    max = _N_BASE_STEPS * 2**k
    j = np.arange(max+1) if k == 0 else np.arange(1, max+1)
    jh = j * h
    pi_2 = np.pi / 2
    u1 = pi_2*np.cosh(jh)
    u2 = pi_2*np.sinh(jh)
    wj = u1 / np.cosh(u2)**2
    xjc = 1 / (np.exp(u2) * np.cosh(u2))
    wj[0] = wj[0] / 2 if k == 0 else wj[0]
    return xjc, wj

def _transform_to_limits(xjc, wj, a, b):
    alpha = (b - a) / 2
    xj = np.concatenate((-alpha * xjc + b, alpha * xjc + a), axis=-1)
    wj = wj*alpha
    wj = np.concatenate((wj, wj), axis=-1)
    invalid = (xj <= a) | (xj >= b)
    wj[invalid] = 0
    return xj, wj

# simple fixed-step integration function
def integrate(func, a, b, args):
    k = 10  # increase this to improve accuracy 
    step0 = _get_base_step()
    step = step0 / 2**k
    xjc, wj = _compute_pair(k, step0)
    xj, wj = _transform_to_limits(xjc, wj, a, b)
    fj = integrand(xj, *args)
    return fj @ wj * step

# your code
def QSzz_1(s,h,z, upsilon, E):
    # exponential form of the QSzz^-1 function to prevent float overflow,
    # also assumes nu==0.5
    numerator = (1+2*s**2*h**2)*np.exp(-2*s*z) + 0.5*(1+np.exp(-4*s*z))
    denominator = 0.5*(1-np.exp(-4*s*z)) - 2*s*h*np.exp(-2*s*z)
    return (3/(2*s*E)) / (numerator/denominator + (3*upsilon/(2*E))*s)

def integrand(s, r, R, upsilon, E, h):
    return s*(R*j0(s*R) - 2*j1(s*R)/s) * QSzz_1(s, h, h, upsilon, E) * j0(s*r)

h = 0.00005
G = 1000
E = 3*G
R= 0.0002
upsilon = 0.06
gamma = 0.072
x_ax = np.linspace(0, 0.0004, 101, endpoint=True, dtype=np.float64)
r, gamma, R, upsilon, E, h = x_ax, gamma, 0.0002, upsilon, E, h
a = 0
b = 1e7
args = (r[:, np.newaxis], R, upsilon, E, h)

%timeit integrate(integrand, a, b, args)
# 2.14 ms ± 82.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

res = integrate(integrand, a, b, args)

stats.gmean(abs((res - res0[0])/res0[0]))
# 0.00011028421945971745

You can trade time for accuracy by increasing k. If you don't want to make the upper limit of integration finite, you can also transform the improper integral into one with finite limits of integration with a variable substitution (but you'll need to increase k to compensate).


Update: style_exact, fit_func, and the call to curve_fit might look like:

def style_exact(r, gamma, R, upsilon, E, h):
    int_out = integrate(integrand, 0, 1e7, args=(r,R,upsilon,E,h))
    return gamma * int_out

def fit_func(x, upsilon, E,h):
    # Ensure that `r` is a CuPy array align led along new axis; the
    # last axis will be used to evaluate integrand at many abscissae.
    x = cp.asarray(x)[:, np.newaxis]
    res = style_exact(x, gamma, 0.0002, upsilon, E, h)
    return cp.asnumpy(res)

popt_hemi, pcov_hemi = curve_fit(fit_func, points_x, points_y, p0=(upsilon, E, h), bounds=([0,0,0],[np.inf, np.inf, np.inf]))
like image 185
Matt Haberland Avatar answered Oct 25 '25 08:10

Matt Haberland



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!