feat: implement NumassBetaWithBump model

2026-02-04 11:20:00 +03:00
parent 45eadb8fc1
commit 3172725643
4 changed files with 78 additions and 35 deletions
--- a/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/NumassBetaWithBump.kt
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/NumassBetaWithBump.kt
@@ -23,7 +23,7 @@ import space.kscience.kmath.expressions.symbol
        val args1 = arguments

        // Второй вызов — переопределяем только mnu2 на B
-        val bVal = arguments.getValue(B)
+        val bVal = arguments.getValue(bumpShift)
        val args2: Map<Symbol, Double> = arguments.toMutableMap().apply {
            put(NumassBeta.mnu2, bVal)
        }
@@ -31,7 +31,7 @@ import space.kscience.kmath.expressions.symbol
        val term1 = NumassBeta(fs, eIn, args1)
        val term2 = NumassBeta(fs, eIn, args2)

-        return term1 + arguments.getValue(A) * term2
+        return term1 + arguments.getValue(bumpAmp) * term2
    }

    override fun derivativeOrNull(symbols: List<Symbol>): Kernel? = when (symbols.size) {
@@ -39,18 +39,18 @@ import space.kscience.kmath.expressions.symbol
        1 -> {
            val sym = symbols.first()
            Kernel { fs, eIn, arguments ->
-                val aVal = arguments.getValue(A)
-                val bVal = arguments.getValue(B)
+                val aVal = arguments.getValue(bumpAmp)
+                val bVal = arguments.getValue(bumpShift)

                val args1 = arguments
                val args2 = arguments.toMutableMap().apply { this[NumassBeta.mnu2] = bVal }

                when (sym) {
                    // ∂/∂A = второй спектр
-                    A -> NumassBeta(fs, eIn, args2)
+                    bumpAmp -> NumassBeta(fs, eIn, args2)

                    // ∂/∂B = A · (∂NumassBeta/∂mnu2) при mnu2 = B
-                    B -> {
+                    bumpShift -> {
                        val derivMnu2 = NumassBeta.derivativeOrNull(listOf(NumassBeta.mnu2))!!
                        aVal * derivMnu2(fs, eIn, args2)
                    }
@@ -88,9 +88,9 @@ import space.kscience.kmath.expressions.symbol

    companion object {
        /** Амплитуда дополнительного компонента */
-        val A: Symbol by symbol
+        val bumpAmp: Symbol by symbol

        /** Фиксированное значение mnu² для дополнительного компонента */
-        val B: Symbol by symbol
+        val bumpShift: Symbol by symbol
    }
 }
--- a/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/beta-compare.kt
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/beta-compare.kt
@@ -1,8 +1,7 @@
 package ru.inr.mass.scripts

-import ru.inr.mass.scripts.NumassBetaWithBump.Companion.A
-import ru.inr.mass.scripts.NumassBetaWithBump.Companion.B
-import space.kscience.plotly.models.AxisType
+import ru.inr.mass.scripts.NumassBetaWithBump.Companion.bumpAmp
+import ru.inr.mass.scripts.NumassBetaWithBump.Companion.bumpShift
 import kotlinx.html.br
 import kotlinx.html.p
 import ru.inr.mass.models.*
@@ -21,7 +20,7 @@ import space.kscience.kmath.expressions.Symbol
 import space.kscience.kmath.real.step
 import space.kscience.plotly.*
 import space.kscience.plotly.models.ScatterMode
-import kotlin.io.path.Path
+import kotlin.math.pow

@UnstableKMathAPI
 fun main() {
@@ -44,8 +43,8 @@ fun main() {
    )

    val argsBump = args.toMutableMap()
-    argsBump[A] = 7e-2
-    argsBump[B] = 10.4
+    argsBump[bumpAmp] = 7e-3
+    argsBump[bumpShift] = (5.3 * 1000.0).pow(2.0)

    val spectrum = constructSpectrum(trapModel, wallModel, adiabaticityModel)
    val spectrumWithBump = constructSpectrum(trapModel, wallModel, adiabaticityModel, NumassBetaWithBump())
@@ -64,9 +63,17 @@ fun main() {
                x.buffer = range
                y.numbers = x.doubles.map { spectrumWithBump(it, argsBump) }
            }
+            scatter {
+                name = "difference (bump - original)"
+                mode = ScatterMode.lines
+                x.buffer = range
+                y.numbers = x.doubles.map {
+                    spectrumWithBump(it, argsBump) - spectrum(it, args)
+                }
+            }
            layout {
                title = "Rear Wall Contribution"
-                yaxis.type = AxisType.log
+//                yaxis.type = AxisType.log
            }
        }
        p {
@@ -81,14 +88,15 @@ fun main() {
            +"wall model = $wallModel"
        }
        br()
-    }.makeFile(Path("./beta-comparison.html"))
+    }.makeFile(show = true)

-    println("HV\toriginal\twith bump")
+    var out = String()
+    out += "HV\toriginal\twith bump\n"
    range.iterator().forEach {
-        println(
-            "${it}\t${spectrum(it, args)}\t${spectrumWithBump(it, argsBump)}"
-        )
+        out += "${it}\t${spectrum(it, args)}\t" +
+        "${spectrumWithBump(it, argsBump)}\n"
    }
+    println(out)
 }

 private fun constructSpectrum(
--- a/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit-custom.kt
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit-custom.kt
@@ -65,6 +65,9 @@ private class FitCustom : CliktCommand(name = "fit-custom") {
    val wallFunc by option().enum<WallModel>().default(WallModel.TSV_2025_12_01).help("wall model")
    val fixRwall by option().flag().help("Don't fit rwall")

+    val bumpShift by option().double().help("bump shift (mnu2) in eV^2")
+    val bumpAmp by option().double().help("bump amplitude")
+
    val adiabacityFunc by option().enum<AdiabaticityModel>().default(AdiabaticityModel.`2024_11`)
        .help("adiabacity tail model")

@@ -73,7 +76,7 @@ private class FitCustom : CliktCommand(name = "fit-custom") {
        val data = parse(spectrum)
        val testData = full?.let { parse(it) } ?: data

-        val initialParams = mapOf(
+        val initialParams = mutableMapOf(
            NBkgSpectrum.norm to norm,
            NBkgSpectrum.bkg to bkg,
            NumassBeta.mnu2 to mnu2,
@@ -85,8 +88,26 @@ private class FitCustom : CliktCommand(name = "fit-custom") {
            rearWall to rwall
        )

+        val source = when {
+            bumpShift == null && bumpAmp == null -> {
+                NumassBeta
+            }
+
+            bumpShift != null && bumpAmp != null -> {
+                initialParams[NumassBetaWithBump.bumpAmp] = bumpAmp!!
+                initialParams[NumassBetaWithBump.bumpShift] = bumpShift!!
+                NumassBetaWithBump()
+            }
+
+            else -> error("bumpShift and bumpAmp must both be set or both be null")
+        }
+
        val spectrumModel = SterileNeutrinoSpectrumRWall(
-            wallModel = wallFunc, trapModel = trapFunc, adiabaticityModel = adiabacityFunc, parallel = true
+            wallModel = wallFunc,
+            trapModel = trapFunc,
+            adiabaticityModel = adiabacityFunc,
+            parallel = true,
+            source = source
        ).withNBkg()

        val fitVars = mutableListOf(NBkgSpectrum.norm)
@@ -104,11 +125,7 @@ private class FitCustom : CliktCommand(name = "fit-custom") {
        }

        val fit = SterileNeutrinoSpectrumRWall.fit(
-            spectrumModel,
-            data,
-            fitVars,
-            initialParams,
-            logger
+            spectrumModel, data, fitVars, initialParams, logger
        )

        val dataPath = Path(spectrum)
--- a/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/grid-fit.kt
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/grid-fit.kt
@@ -16,7 +16,6 @@ import kotlinx.coroutines.runBlocking
 import ru.inr.mass.models.*
 import ru.inr.mass.scripts.SterileNeutrinoSpectrumRWall.Companion.rearWall
 import space.kscience.kmath.UnstableKMathAPI
-import space.kscience.kmath.expressions.Symbol
 import space.kscience.kmath.misc.Loggable
 import space.kscience.kmath.optimization.*
 import space.kscience.plotly.*
@@ -80,11 +79,14 @@ private class GridArgs : CliktCommand() {
    val extraM: List<Double> by option("--extra-m").double().multiple().help("Extra mass points (keV), can be repeated")
    val extraU2: List<Double> by option("--extra-u2").double().multiple().help("Extra u² points, can be repeated")

+    val bumpShift by option().double().help("bump shift (mnu2) in eV^2")
+    val bumpAmp by option().double().help("bump amplitude")
+
    @UnstableKMathAPI
    override fun run() {
        val postfix = if (postfix != null) "-$postfix" else ""

-        val fitParamsBase: Map<Symbol, Double> = mapOf(
+        val fitParamsBase = mutableMapOf(
            NBkgSpectrum.norm to norm,
            NBkgSpectrum.bkg to bkg,
            NumassBeta.mnu2 to mnu2,
@@ -94,8 +96,26 @@ private class GridArgs : CliktCommand() {
            rearWall to rwall
        )

+        val source = when {
+            bumpShift == null && bumpAmp == null -> {
+                NumassBeta
+            }
+
+            bumpShift != null && bumpAmp != null -> {
+                fitParamsBase[NumassBetaWithBump.bumpAmp] = bumpAmp!!
+                fitParamsBase[NumassBetaWithBump.bumpShift] = bumpShift!!
+                NumassBetaWithBump()
+            }
+
+            else -> error("bumpShift and bumpAmp must both be set or both be null")
+        }
+
        val spectrumModel: NBkgSpectrum = SterileNeutrinoSpectrumRWall(
-            wallModel = wallFunc, trapModel = trapFunc, adiabaticityModel = adiabacityFunc, parallel = true
+            wallModel = wallFunc,
+            trapModel = trapFunc,
+            adiabaticityModel = adiabacityFunc,
+            parallel = true,
+            source = source
        ).withNBkg()

        val fitVars = mutableListOf(NBkgSpectrum.norm)
@@ -104,7 +124,7 @@ private class GridArgs : CliktCommand() {
        if (!fixTrap) fitVars += NumassTransmission.trap
        if (!fixRwall) fitVars += rearWall

-        val data = ru.inr.mass.scripts.parse(spectrum)
+        val data = parse(spectrum)
        val dataPath = Path(spectrum)
        val dataForView = if (full != null) parse(full!!) else data

@@ -124,8 +144,7 @@ private class GridArgs : CliktCommand() {
            val chi2Matrix = mutableMapOf<Pair<Double, Double>, Double?>()

            val tsvPath = Path(
-                dataPath.parent.toString(),
-                dataPath.nameWithoutExtension + postfix + ".chi2_grid.tsv"
+                dataPath.parent.toString(), dataPath.nameWithoutExtension + postfix + ".chi2_grid.tsv"
            )
            val tsvFile = File(tsvPath.toString())

@@ -250,8 +269,7 @@ private fun saveChi2GridToTsv(
 ) {
    val dataPath = Path(spectrumFile)
    val tsvPath = Path(
-        dataPath.parent.toString(),
-        dataPath.nameWithoutExtension + postfix + ".chi2_grid.tsv"
+        dataPath.parent.toString(), dataPath.nameWithoutExtension + postfix + ".chi2_grid.tsv"
    )

    File(tsvPath.toString()).printWriter().use { out ->