NPM/numass
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[enh] make SterileNeutrinoSpectrumFast with parallel integration

Browse Source
This commit is contained in:
Vasily Chernov
2025-03-14 13:25:20 +03:00
parent cb720ca89e
commit 1373bd73bd
4 changed files with 241 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

238
numass-model/src/jvmMain/kotlin/ru/inr/mass/models/SterileNeutrinoSpectrumFast.kt Normal file
View File

@@ -0,0 +1,238 @@
package ru.inr.mass.models
import kotlinx.coroutines.*
import ru.inr.mass.models.NumassBeta.e0
import ru.inr.mass.models.NumassBeta.mnu2
import ru.inr.mass.models.NumassBeta.msterile2
import ru.inr.mass.models.NumassBeta.u2
import ru.inr.mass.models.NumassTransmission.Companion.thickness
import ru.inr.mass.models.NumassTransmission.Companion.trap
import space.kscience.kmath.expressions.Symbol
import space.kscience.kmath.expressions.derivative
import kotlin.math.max
/**
* @param source variables:Eo offset,Ein; parameters: "mnu2", "msterile2", "U2"
* @param transmission variables:Ein,Eout; parameters: "A"
* @param resolution variables:Eout,U; parameters: "X", "trap"
*/
public open class SterileNeutrinoSpectrumFast(
public val source: DifferentiableKernel = NumassBeta,
public val transmission: DifferentiableKernel = NumassTransmission(),
public val resolution: DifferentiableKernel = NumassResolution(),
public val fss: FSS? = null,
) : DifferentiableSpectrum {
/**
* auxiliary function for trans-res convolution
*/
private val transRes: DifferentiableKernel = TransRes()
// private boolean useMC;
//private val fast: Boolean = configuration.getBoolean("fast", true)
override fun invoke(u: Double, arguments: Map<Symbol, Double>): Double {
return convolute(u, source, transRes, arguments)
}
override fun derivativeOrNull(symbols: List<Symbol>): Spectrum? {
if (symbols.isEmpty()) return this
return when (symbols.singleOrNull() ?: TODO("First derivatives only")) {
u2, msterile2, mnu2, e0 -> Spectrum { u, arguments ->
convolute(u, source.derivative(symbols), transRes, arguments)
}
thickness, trap -> Spectrum { u, arguments ->
convolute(u, source, transRes.derivative(symbols), arguments)
}
else -> null
}
}
/**
* Direct Gauss-Legendre integration
*
* @param u
* @param sourceFunction
* @param transResFunction
* @param arguments
* @return
*/
private fun convolute(
u: Double,
sourceFunction: Kernel,
transResFunction: Kernel,
arguments: Map<Symbol, Double>,
): Double {
val eMax = arguments.getValue(e0) + 5.0
if (u >= eMax) {
return 0.0
}
val xs = createRange(u, eMax)
return runBlocking(Dispatchers.Default) {
parallelIntegrate(xs) { eIn ->
sumByFSS(eIn, sourceFunction, arguments) * transResFunction(eIn, u, arguments)
}
}
}
private fun sumByFSS(eIn: Double, sourceFunction: Kernel, arguments: Map<Symbol, Double>): Double {
return if (fss == null) {
sourceFunction(0.0, eIn, arguments)
} else {
(0 until fss.size).sumOf { fss.ps[it] * sourceFunction(fss.es[it], eIn, arguments) }
}
}
private inner class TransRes : DifferentiableKernel {
override fun invoke(eIn: Double, u: Double, arguments: Map<Symbol, Double>): Double {
val p0 = NumassTransmission.p0(eIn, arguments)
return p0 * resolution(eIn, u, arguments) + lossRes(transmission, eIn, u, arguments)
}
override fun derivativeOrNull(symbols: List<Symbol>): Kernel? {
if (symbols.isEmpty()) return this
return when (symbols.singleOrNull() ?: TODO("First derivatives only")) {
thickness -> null//TODO implement p0 derivative
trap -> Kernel { eIn, u, arguments -> lossRes(transmission.derivative(symbols), eIn, u, arguments) }
else -> null
}
}
private fun lossRes(
transFunc: Kernel,
eIn: Double,
u: Double,
arguments: Map<Symbol, Double>,
): Double {
val xs = createRange(u, eIn)
return runBlocking {
parallelIntegrate(xs) { eOut: Double ->
transFunc(eIn, eOut, arguments) * resolution(eOut, u, arguments)
}
}
}
}
public companion object
}
/**
* Создание границ, с шагами такими же, как в оригинальной функции [SterileNeutrinoSpectrum.convolute]
*/
private fun createRange(x1: Double, x2: Double): DoubleArray {
val result = mutableListOf<Double>()
var currentValue = x1
val steps = listOf(2.0, 5.0, 8.0, 15.0, 20.0)
var stepIndex = 0
while (currentValue <= x2) {
result.add(currentValue)
currentValue += if (stepIndex < steps.size) {
steps[stepIndex++]
} else {
25.0
}
}
// Ensure the last element is exactly x2 if possible
if (result.isNotEmpty() && result.last() != x2) {
if (result.last() < x2) {
result.add(x2)
} else {
result[result.lastIndex] = x2
}
}
return result.toDoubleArray()
}
/**
* Параллельный интегратор одномерной функции методом трапеций.
*
* @param x Сетка координат.
* @param f Функция, которую нужно интегрировать.
* @return Значение интеграла.
*/
private suspend fun parallelIntegrate(
x: DoubleArray,
f: (Double) -> Double,
): Double = coroutineScope {
val numThreads = Runtime.getRuntime().availableProcessors()
val chunkSize = max(9, x.size / numThreads)
val results = mutableListOf<Deferred<Double>>()
x.asSequence().windowed(chunkSize+1, chunkSize, partialWindows=true).forEach {
results.add(async {
gaussQuadrature(it.toDoubleArray(), f)
})
}
val partialResults = results.awaitAll()
return@coroutineScope partialResults.sum()
}
/**
* Интегрирование одномерной функции методом трапеций.
*
* @param x Сетка координат.
* @param f Функция, которую нужно интегрировать.
* @return Значение интеграла или null, если функция не определена в какой-то точке.
*/
private fun integrateTrapezoidal(x: DoubleArray, f: (Double) -> Double): Double {
if (x.size < 2) {
return 0.0 // Так как parallelIntegrate передает куски внахлест, то остаток размера 1 дает 0,
// Так как он уже был просуммирован в предыдущем куске
}
var integral = 0.0
for (i in 0 until x.size - 1) {
val y1 = f(x[i])
val y2 = f(x[i + 1])
integral += (y1 + y2) * (x[i + 1] - x[i]) / 2.0
}
return integral
}
// Коэффициенты и узлы Гаусса-Лежандра для 5 точек
private val gaussPoints = doubleArrayOf(
-0.9061798459,
-0.5384693101,
0.0,
0.5384693101,
0.9061798459
)
private val gaussWeights = doubleArrayOf(
0.2369268851,
0.4786286705,
0.5688888889,
0.4786286705,
0.2369268851
)
private fun gaussQuadrature(grid: DoubleArray, func: (Double) -> Double): Double {
var integral = 0.0
for (i in 0 until grid.size - 1) {
val a = grid[i]
val b = grid[i + 1]
val h = (b - a) / 2.0
val mid = (a + b) / 2.0
for (j in gaussPoints.indices) {
val x = mid + h * gaussPoints[j]
integral += h * gaussWeights[j] * func(x)
}
}
return integral
}

4
numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit-custom.kt
View File

@@ -38,7 +38,7 @@ class CustomSterileNeutrinoSpectrum(
transmission: DifferentiableKernel = NumassTransmission(),
resolution: DifferentiableKernel = NumassResolution(),
fss: FSS? = null
): SterileNeutrinoSpectrum(source, transmission, resolution, fss) {
): SterileNeutrinoSpectrumFast(source, transmission, resolution, fss) {
override fun invoke(u: Double, arguments: Map<Symbol, Double>): Double {
val rearWall = arguments[rearWall]!!
@@ -168,7 +168,7 @@ suspend fun processCustom(
},
adjustX = false,
),
resolution = NumassResolution(1.7e-4, tailFunction = {e,u -> adiab2024(e,u)})
resolution = NumassResolution(1.7e-4, tailFunction = {e,u -> adiabacity(e,u) })
).withNBkg()

2
numass-workspace/src/main/kotlin/ru/inr/mass/scripts/fit.kt
View File

@@ -160,7 +160,7 @@ private suspend fun process(spectrumFile: String, full: String?, postfix: String
},
adjustX = false,
),
resolution = NumassResolution(1.7e-4, tailFunction = {e,u -> adiab2024(e,u)})
resolution = NumassResolution(1.7e-4, tailFunction = {e,u -> adiabacity(e,u) })
).withNBkg()

4
numass-workspace/src/main/kotlin/ru/inr/mass/scripts/get-spectrum.kt
View File

@@ -13,13 +13,9 @@ import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.runBlocking
import ru.inr.mass.models.*
import space.kscience.kmath.data.XYErrorColumnarData
import space.kscience.kmath.UnstableKMathAPI
import space.kscience.kmath.expressions.Symbol
import space.kscience.kmath.operations.toTypedArray
import space.kscience.kmath.real.plus
import space.kscience.kmath.real.step
import space.kscience.plotly.models.HoverMode
import kotlin.math.pow
private class CalcSpectrum : CliktCommand() {