From cf02b2d8bac3a706b7e1ed1efced7edb4f2bf930 Mon Sep 17 00:00:00 2001
From: darksnake <altavir@gmail.com>
Date: Tue, 24 Jan 2023 13:05:38 +0300
Subject: [PATCH] Add fit demo (not working yet)

---
 .../main/kotlin/ru/inr/mass/scripts/fit.kt    | 72 +++++++++++++++++++
 .../mass/workspace/{fit.kt => dataUtils.kt}   | 31 ++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit.kt
 rename numass-workspace/src/main/kotlin/ru/inr/mass/workspace/{fit.kt => dataUtils.kt} (55%)

diff --git a/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit.kt b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit.kt
new file mode 100644
index 0000000..17cf8d6
--- /dev/null
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit.kt
@@ -0,0 +1,72 @@
+package ru.inr.mass.scripts
+
+import ru.inr.mass.models.*
+import ru.inr.mass.workspace.buffer
+import ru.inr.mass.workspace.fitWith
+import ru.inr.mass.workspace.generate
+import space.kscience.kmath.expressions.Symbol
+import space.kscience.kmath.misc.UnstableKMathAPI
+import space.kscience.kmath.operations.asSequence
+import space.kscience.kmath.optimization.*
+import space.kscience.kmath.real.step
+import space.kscience.plotly.*
+import space.kscience.plotly.models.ScatterMode
+import kotlin.math.pow
+
+
+@OptIn(UnstableKMathAPI::class)
+suspend fun main() {
+    val spectrum: NBkgSpectrum = SterileNeutrinoSpectrum(
+        fss = FSS.default,
+//        resolution = NumassResolution(8.2e-5)
+    ).withNBkg()
+
+    val args: Map<Symbol, Double> = mapOf(
+        NBkgSpectrum.norm to 8e5,
+        NBkgSpectrum.bkg to 2.0,
+        NumassBeta.mnu2 to 0.0,
+        NumassBeta.e0 to 18575.0,
+        NumassBeta.msterile2 to 1000.0.pow(2),
+        NumassBeta.u2 to 1e-2,
+        NumassTransmission.thickness to 0.1,
+        NumassTransmission.trap to 1.0
+    )
+
+    val timePerPoint = 30.0
+
+    val strategy = (12000.0..19000.0 step 100.0).asSequence().associateWith { timePerPoint }
+
+    val generatedData = spectrum.generate(strategy, args)
+
+    val fit: XYFit = generatedData.fitWith(
+        optimizer = QowOptimizer,
+        modelExpression = spectrum,
+        startingPoint = args,
+        OptimizationParameters(NBkgSpectrum.norm, NBkgSpectrum.bkg, NumassBeta.e0)
+    )
+
+    val fitResult = fit.resultPoint
+
+    Plotly.plot {
+        scatter {
+            name = "Generated"
+            mode = ScatterMode.markers
+            x.buffer = generatedData.x
+            y.buffer = generatedData.y
+        }
+
+        scatter {
+            name = "Initial"
+            mode = ScatterMode.lines
+            x.buffer = 12000.0..18600.0 step 50.0
+            y.numbers = x.doubles.map { spectrum(it, args) }
+        }
+
+        scatter {
+            name = "Fit"
+            mode = ScatterMode.lines
+            x.buffer = 12000.0..18600.0 step 10.0
+            y.numbers = x.doubles.map { spectrum(it, fitResult) }
+        }
+    }.makeFile()
+}
\ No newline at end of file
diff --git a/numass-workspace/src/main/kotlin/ru/inr/mass/workspace/fit.kt b/numass-workspace/src/main/kotlin/ru/inr/mass/workspace/dataUtils.kt
similarity index 55%
rename from numass-workspace/src/main/kotlin/ru/inr/mass/workspace/fit.kt
rename to numass-workspace/src/main/kotlin/ru/inr/mass/workspace/dataUtils.kt
index 4f43895..7e48c6a 100644
--- a/numass-workspace/src/main/kotlin/ru/inr/mass/workspace/fit.kt
+++ b/numass-workspace/src/main/kotlin/ru/inr/mass/workspace/dataUtils.kt
@@ -2,13 +2,20 @@
 
 package ru.inr.mass.workspace
 
+import ru.inr.mass.models.Spectrum
 import space.kscience.kmath.data.XYColumnarData
+import space.kscience.kmath.data.XYErrorColumnarData
 import space.kscience.kmath.expressions.DifferentiableExpression
 import space.kscience.kmath.expressions.Symbol
 import space.kscience.kmath.misc.FeatureSet
 import space.kscience.kmath.misc.Loggable
 import space.kscience.kmath.misc.UnstableKMathAPI
 import space.kscience.kmath.optimization.*
+import space.kscience.kmath.samplers.PoissonSampler
+import space.kscience.kmath.stat.RandomGenerator
+import space.kscience.kmath.stat.next
+import space.kscience.kmath.structures.asBuffer
+import kotlin.math.sqrt
 
 public suspend fun XYColumnarData<Double, Double, Double>.fitWith(
     optimizer: Optimizer<Double, XYFit>,
@@ -33,4 +40,28 @@ public suspend fun XYColumnarData<Double, Double, Double>.fitWith(
         xSymbol
     )
     return optimizer.optimize(problem)
+}
+
+public suspend fun Spectrum.generate(
+    strategy: Map<Double, Double>,
+    arguments: Map<Symbol, Double>,
+    generator: RandomGenerator = RandomGenerator.default,
+): XYErrorColumnarData<Double, Double, Double> {
+    val xs = mutableListOf<Double>()
+    val ys = mutableListOf<Double>()
+    val errors = mutableListOf<Double>()
+
+    strategy.forEach { (x, time) ->
+        xs.add(x)
+        val mu = invoke(x, arguments) * time
+        val y = PoissonSampler(mu).next(generator) / time
+        val error = sqrt(mu) / time
+        ys.add(y)
+        errors.add(error)
+    }
+    return XYErrorColumnarData.of(
+        xs.toTypedArray().asBuffer(),
+        ys.toTypedArray().asBuffer(),
+        errors.toTypedArray().asBuffer()
+    )
 }
\ No newline at end of file