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]
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]))