""" JAX port of the MSHT20 parametrization, transcribed 1:1 from FPPDF (github.com/FPPDF/fppdf, src/fppdf/{pdfs.py, chebyshevs.py}, main @ 2026-07-07). Every function mirrors its FPPDF counterpart by name; parity is enforced by analysis/msht20/parity_test.py against the FPPDF implementation itself. Conventions (FPPDF defaults, = basis_pars with Cheb_8=False, g_second_term=True, g_cheb7=False, asp_fix=True, dvd_eq_uvd=False, t8_int=False): 73-slot parameter vector: uv [ 0: 9] A_uv, del_uv, eta_uv, c1..c6 (A_uv <- valence sum rule) dv [ 9:18] A_dv, del_dv, eta_dv, c1..c6 (A_dv <- valence sum rule) sea [18:27] A_S, del_S, eta_S, c1..c6 s+ [27:36] A_sp, del_sp, eta_sp, c1..c6 (del_sp <- del_S if asp_fix) g [36:46] A_g, eta_g+, del_g+, cg1..cg4, A_g-, eta_g-, del_g- (A_g <- momentum sum rule) s- [46:56] A_sm, del_sm, eta_sm, x0, c1..c6 (x0 <- strange number sum rule) db-ub [56:64] A_rho, eta_rho, c1..c6 (note: NO delta slot) c+ [64:73] A_c, del_c, eta_c, c1..c6 (zero for perturbative charm) All PDFs are x*f(x). Flavour dispatch pdfs_msht(ipdf,...) uses LHAPDF-like ids: 0=g, 1=d, 2=u, 3=s, -1=dbar, -2=ubar, -3=sbar, +-4=c (=c+/2). """ import jax import jax.numpy as jnp from jax.scipy.special import gammaln jax.config.update("jax_enable_x64", True) # indices (Cheb_8=False layout) I_UV, I_DV, I_SEA, I_SP = slice(0, 9), slice(9, 18), slice(18, 27), slice(27, 36) I_G, I_SM, I_DBUB, I_CH = slice(36, 46), slice(46, 56), slice(56, 64), slice(64, 73) NPARS = 73 def _gamma(z): """Signed gamma via gammaln + gammasgn (math.gamma semantics for z<0).""" try: from jax.scipy.special import gammasgn return gammasgn(z) * jnp.exp(gammaln(z)) except ImportError: # older jax: reflection-based sign sign = jnp.where(z > 0, 1.0, jnp.sign(jnp.sin(jnp.pi * z)) * jnp.where(jnp.sin(jnp.pi * z) == 0, 0.0, 1.0)) return sign * jnp.exp(gammaln(z)) def I_beta(a, b): """FPPDF chebyshevs.I — Euler beta with Stirling branch for b >= 100.""" small = _gamma(a + 1.0) * _gamma(b + 1.0) / _gamma(a + b + 2.0) big = (jnp.sqrt(b / (a + b + 1.0)) * jnp.power(b / (a + b + 1.0), b) * jnp.power(a + b + 1.0, -a - 1.0) * _gamma(a + 1.0) * jnp.exp(a + 1.0)) return jnp.where(b < 100.0, small, big) # ∫ x^a (1-x)^b y^k dx with y = 1 - 2 sqrt(x): binomial expansions (FPPDF Iy1..Iy8) def Iy1(a, b): return I_beta(a, b) - 2.0 * I_beta(a + 0.5, b) def Iy2(a, b): return I_beta(a, b) - 4.0 * I_beta(a + 0.5, b) + 4.0 * I_beta(a + 1.0, b) def Iy3(a, b): return (I_beta(a, b) - 6.0 * I_beta(a + 0.5, b) + 12.0 * I_beta(a + 1.0, b) - 8.0 * I_beta(a + 1.5, b)) def Iy4(a, b): return (I_beta(a, b) - 8.0 * I_beta(a + 0.5, b) + 24.0 * I_beta(a + 1.0, b) - 32.0 * I_beta(a + 1.5, b) + 16.0 * I_beta(a + 2.0, b)) def Iy5(a, b): return (I_beta(a, b) - 10.0 * I_beta(a + 0.5, b) + 40.0 * I_beta(a + 1.0, b) - 80.0 * I_beta(a + 1.5, b) + 80.0 * I_beta(a + 2.0, b) - 32.0 * I_beta(a + 2.5, b)) def Iy6(a, b): return (I_beta(a, b) - 12.0 * I_beta(a + 0.5, b) + 60.0 * I_beta(a + 1.0, b) - 160.0 * I_beta(a + 1.5, b) + 240.0 * I_beta(a + 2.0, b) - 192.0 * I_beta(a + 2.5, b) + 64.0 * I_beta(a + 3.0, b)) def Iy7(a, b): return (I_beta(a, b) - 14.0 * I_beta(a + 0.5, b) + 84.0 * I_beta(a + 1.0, b) - 280.0 * I_beta(a + 1.5, b) + 560.0 * I_beta(a + 2.0, b) - 672.0 * I_beta(a + 2.5, b) + 448.0 * I_beta(a + 3.0, b) - 128.0 * I_beta(a + 3.5, b)) def Iy8(a, b): return (I_beta(a, b) - 16.0 * I_beta(a + 0.5, b) + 112.0 * I_beta(a + 1.0, b) - 448.0 * I_beta(a + 1.5, b) + 1120.0 * I_beta(a + 2.0, b) - 1792.0 * I_beta(a + 2.5, b) + 1792.0 * I_beta(a + 3.0, b) - 1024.0 * I_beta(a + 3.5, b) + 256.0 * I_beta(a + 4.0, b)) # ∫ x^a (1-x)^b T_k(y) dx — Chebyshev expansions in Iy (FPPDF Ic1..Ic8) def Ic1(a, b): return Iy1(a, b) def Ic2(a, b): return 2.0 * Iy2(a, b) - I_beta(a, b) def Ic3(a, b): return 4.0 * Iy3(a, b) - 3.0 * Iy1(a, b) def Ic4(a, b): return 8.0 * Iy4(a, b) - 8.0 * Iy2(a, b) + I_beta(a, b) def Ic5(a, b): return 16.0 * Iy5(a, b) - 20.0 * Iy3(a, b) + 5.0 * Iy1(a, b) def Ic6(a, b): return 32.0 * Iy6(a, b) - 48.0 * Iy4(a, b) + 18.0 * Iy2(a, b) - I_beta(a, b) def Ic7(a, b): return 64.0 * Iy7(a, b) - 112.0 * Iy5(a, b) + 56.0 * Iy3(a, b) - 7.0 * Iy1(a, b) def Ic8(a, b): return (128.0 * Iy8(a, b) - 256.0 * Iy6(a, b) + 160.0 * Iy4(a, b) - 32.0 * Iy2(a, b) + I_beta(a, b)) _IC = [None, Ic1, Ic2, Ic3, Ic4, Ic5, Ic6, Ic7, Ic8] def cheb_msht(i, x): """T_i(y), y = 1 - 2 sqrt(x). Explicit polynomials, as in FPPDF.""" y = 1.0 - 2.0 * jnp.sqrt(x) if i == 1: return y if i == 2: return 2.0 * y**2 - 1.0 if i == 3: return 4.0 * y**3 - 3.0 * y if i == 4: return 8.0 * y**4 - 8.0 * y**2 + 1.0 if i == 5: return 16.0 * y**5 - 20.0 * y**3 + 5.0 * y if i == 6: return 32.0 * y**6 - 48.0 * y**4 + 18.0 * y**2 - 1.0 if i == 7: return 64.0 * y**7 - 112.0 * y**5 + 56.0 * y**3 - 7.0 * y if i == 8: return 128.0 * y**8 - 256.0 * y**6 + 160.0 * y**4 - 32.0 * y**2 + 1.0 raise ValueError(i) def _cheb_series(x, coeffs): """1 + sum_i coeffs[i] * T_{i+1}(y(x)) for 6 coefficients.""" out = 1.0 for i, c in enumerate(coeffs, start=1): out = out + c * cheb_msht(i, x) return out def q_msht(x, ain): """Quark-like shape: A x^del (1-x)^eta (1 + sum c_i T_i). Cheb_8=False.""" aq, delq, etaq = ain[0], ain[1], ain[2] shape = aq * jnp.power(1.0 - x, etaq) * jnp.power(x, delq) * _cheb_series(x, [ain[3+i] for i in range(6)]) return jnp.where(aq < 1e-20, jnp.zeros_like(x) * aq, shape) def dbub_msht(x, ain): """dbar-ubar shape: NO x^delta factor. etaq<0 clips x at 0.999.""" aq, etaq = ain[0], ain[1] xx = jnp.where(etaq < 0.0, jnp.minimum(x, 0.999), x) return aq * jnp.power(1.0 - xx, etaq) * _cheb_series(xx, [ain[2+i] for i in range(6)]) def g_msht(x, ain, two_terms=True): """Gluon: main Chebyshev term (cg1..cg4) + optional negative second term.""" agp, etagp, delgp = ain[0], ain[1], ain[2] agm, etagm, delgm = ain[7], ain[8], ain[9] main = agp * jnp.power(1.0 - x, etagp) * jnp.power(x, delgp) * _cheb_series(x, [ain[3+i] for i in range(4)]) if two_terms: return main + agm * jnp.power(1.0 - x, etagm) * jnp.power(x, delgm) # not used at FPPDF defaults; kept for completeness (etagm/delgm act as c5/c6) y5 = agp * jnp.power(1.0 - x, etagp) * jnp.power(x, delgp) * ( etagm * cheb_msht(5, x) + delgm * cheb_msht(6, x)) return main + y5 def sm_msht(x, ain): """s-minus: A x^del (1-x)^eta (1 - x/x0)(1 + sum c_i T_i).""" asm, delsm, etasm, x0 = ain[0], ain[1], ain[2], ain[3] out = asm * jnp.power(1.0 - x, etasm) * jnp.power(x, delsm) * (1.0 - x / x0) return out * _cheb_series(x, [ain[4+i] for i in range(6)]) # ---------------- sum-rule integrals ---------------- def _cheb_int(a, b, coeffs): """I(a,b) + sum_i coeffs[i] * Ic_{i+1}(a,b) for 6 coefficients.""" out = I_beta(a, b) for i, c in enumerate(coeffs, start=1): out = out + c * _IC[i](a, b) return out def qv_int(aq, iq): """FPPDF qv_int (xmin=0): iq=1 number integral (del-1), iq=2 momentum (del).""" delq = jnp.where(iq == 1, aq[1] - 1.0, aq[1]) return aq[0] * _cheb_int(delq, aq[2], [aq[3+i] for i in range(6)]) def qv_norm(iq, aq): i1 = qv_int(aq, 1) return (2.0 / i1) if iq == 1 else (1.0 / i1) def q_msht_lowx_norm(ain): return ain[0] * (1.0 + sum(ain[3+i] for i in range(6))) def sp_norm_fix(asea, ain): norm_cheb = 1.0 + sum(ain[3+i] for i in range(6)) return asea / norm_cheb / 3.0 def sm1_int(ain): return ain[0] * _cheb_int(ain[1] - 1.0, ain[2], [ain[4+i] for i in range(6)]) def sm2_int(ain): return ain[0] / ain[3] * _cheb_int(ain[1], ain[2], [ain[4+i] for i in range(6)]) def smin_norm(asm): return sm2_int(asm) / sm1_int(asm) def int_g1_msht(ain, two_terms=True): """NB: FPPDF multiplies by ain[0] (A_g slot) — normally 1.0 pre-sum-rule.""" assert two_terms, "only the FPPDF-default two_terms gluon is ported" return ain[0] * _cheb_int(ain[2], ain[1], [ain[3], ain[4], ain[5], ain[6], 0.0, 0.0]) def int_g2_msht(ain, two_terms=True): return ain[7] * I_beta(ain[9], ain[8]) if two_terms else 0.0 def msum_ag(pars): """Momentum sum rule -> gluon normalization A_g.""" outng = (qv_int(pars[I_UV], 2) + qv_int(pars[I_DV], 2) + qv_int(pars[I_SEA], 2) + qv_int(pars[I_CH], 2)) outg1 = int_g1_msht(pars[I_G]) outg2 = int_g2_msht(pars[I_G]) return (1.0 - outng - outg2) / outg1 def sumrules(parin): """Fill the sum-rule slots: A_uv, A_dv, A_g, x0(s-), del_sp tie.""" p = parin # asp_fix: del_sp = del_S p = p.at[27 + 1].set(p[18 + 1]) # s- zero-strangeness x0 p = p.at[46 + 3].set(smin_norm(p[I_SM])) # valence counting p = p.at[0].set(qv_norm(1, p[I_UV])) p = p.at[9].set(qv_norm(2, p[I_DV])) # momentum sum -> gluon A_g p = p.at[36].set(msum_ag(p)) return p # ---------------- flavour assembly (pdfs_msht) ---------------- def _basis(ipdf, pars, x): if ipdf == 1: return q_msht(x, pars[I_UV]) if ipdf == 2: return q_msht(x, pars[I_DV]) if ipdf == 3: return q_msht(x, pars[I_SEA]) if ipdf == 4: return q_msht(x, pars[I_SP]) if ipdf == 5: return g_msht(x, pars[I_G]) if ipdf == 6: return sm_msht(x, pars[I_SM]) if ipdf == 7: return dbub_msht(x, pars[I_DBUB]) if ipdf == 8: return q_msht(x, pars[I_CH]) raise ValueError(ipdf) def pdfs_msht(ipdf, pars, x): """LHAPDF-id flavours from the MSHT basis (mirrors FPPDF pdfs_msht).""" etaq = pars[57] etaqt = jnp.power(1.0 - x, -etaq) valid = jnp.where(etaq < 0.0, etaqt > 1e-20, jnp.ones_like(x, dtype=bool)) dbdub = jnp.where(valid, _basis(7, pars, x), 0.0) sp = _basis(4, pars, x) sm = _basis(6, pars, x) sea = _basis(3, pars, x) dbpub = (sea - sp) / 2.0 if ipdf == -3: return (sp - sm) / 2.0 if ipdf == 3: return (sp + sm) / 2.0 if ipdf == -2: return dbpub / (1.0 + dbdub) * valid if ipdf == -1: return dbpub * jnp.where(valid, dbdub / (1.0 + dbdub), 1.0) if ipdf == 0: return _basis(5, pars, x) if ipdf == 1: db = dbpub * jnp.where(valid, dbdub / (1.0 + dbdub), 1.0) return _basis(2, pars, x) + db if ipdf == 2: ub = valid * dbpub / (1.0 + dbdub) return _basis(1, pars, x) + ub if ipdf in (-4, 4): return _basis(8, pars, x) / 2.0 return jnp.zeros_like(x) # ---------------- starting parameters ---------------- def load_publication_start(dat_path): """Assemble the 73-slot vector from input_mshtfit_publication_shortmin.dat. Sections (free values only; sum-rule slots get placeholders=1.0, charm=0): uv 8, dv 8, sea 9, s+ 9, g 9, s- 9(x0 inserted), db-ub 8.""" import re vals = [] for line in open(dat_path): line = line.strip() m = re.match(r"^([-+0-9.eE]+)\s+[01]\s*$", line) if m: vals.append(float(m.group(1))) assert len(vals) == 60, f"expected 60 values, got {len(vals)}" uv, dv, sea, sp, g, smv, dbub = (vals[0:8], vals[8:16], vals[16:25], vals[25:34], vals[34:43], vals[43:52], vals[52:60]) import numpy as np p = np.zeros(NPARS) p[0], p[1:9] = 1.0, uv p[9], p[10:18] = 1.0, dv p[18:27] = sea p[27:36] = sp p[36], p[37:46] = 1.0, g p[46:49] = smv[0:3] # A_sm, del_sm, eta_sm p[49] = 1.0 # x0 placeholder (sum rule) p[50:56] = smv[3:9] # c1..c6 (zeros in publication file) p[56:64] = dbub # charm slots stay 0 (perturbative charm) return jnp.array(p)