Add fit demo (not working yet)

This commit is contained in:
Alexander Nozik 2023-01-24 13:05:38 +03:00
parent d9b4707da1
commit cf02b2d8ba
2 changed files with 103 additions and 0 deletions

View File

@ -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()
}

View File

@ -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>,
@ -34,3 +41,27 @@ public suspend fun XYColumnarData<Double, Double, Double>.fitWith(
)
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()
)
}