diff --git a/examples/src/main/kotlin/kscience/kmath/ast/KotlingradSupport.kt b/examples/src/main/kotlin/kscience/kmath/ast/KotlingradSupport.kt index 5fbb5b86a..5acd97e3d 100644 --- a/examples/src/main/kotlin/kscience/kmath/ast/KotlingradSupport.kt +++ b/examples/src/main/kotlin/kscience/kmath/ast/KotlingradSupport.kt @@ -3,7 +3,6 @@ package kscience.kmath.ast import kscience.kmath.asm.compile import kscience.kmath.expressions.invoke import kscience.kmath.expressions.symbol -import kscience.kmath.kotlingrad.DifferentiableMstExpression import kscience.kmath.kotlingrad.derivative import kscience.kmath.kotlingrad.differentiable import kscience.kmath.operations.RealField diff --git a/kmath-commons/src/main/kotlin/kscience/kmath/commons/expressions/DerivativeStructureExpression.kt b/kmath-commons/src/main/kotlin/kscience/kmath/commons/expressions/DerivativeStructureExpression.kt index 244dc1314..345babe8b 100644 --- a/kmath-commons/src/main/kotlin/kscience/kmath/commons/expressions/DerivativeStructureExpression.kt +++ b/kmath-commons/src/main/kotlin/kscience/kmath/commons/expressions/DerivativeStructureExpression.kt @@ -95,10 +95,10 @@ public class DerivativeStructureField( public override operator fun Number.plus(b: DerivativeStructure): DerivativeStructure = b + this public override operator fun Number.minus(b: DerivativeStructure): DerivativeStructure = b - this - public companion object : AutoDiffProcessor { - override fun process(function: DerivativeStructureField.() -> DerivativeStructure): DifferentiableExpression { - return DerivativeStructureExpression(function) - } + public companion object : + AutoDiffProcessor> { + public override fun process(function: DerivativeStructureField.() -> DerivativeStructure): DifferentiableExpression> = + DerivativeStructureExpression(function) } } @@ -108,7 +108,7 @@ public class DerivativeStructureField( */ public class DerivativeStructureExpression( public val function: DerivativeStructureField.() -> DerivativeStructure, -) : DifferentiableExpression { +) : DifferentiableExpression> { public override operator fun invoke(arguments: Map): Double = DerivativeStructureField(0, arguments).function().value diff --git a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/DifferentiableExpression.kt b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/DifferentiableExpression.kt index ac1f4bc20..890ad5f71 100644 --- a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/DifferentiableExpression.kt +++ b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/DifferentiableExpression.kt @@ -1,29 +1,37 @@ package kscience.kmath.expressions /** - * An expression that provides derivatives + * Represents expression which structure can be differentiated. + * + * @param T the type this expression takes as argument and returns. + * @param R the type of expression this expression can be differentiated to. */ -public interface DifferentiableExpression : Expression { - public fun derivativeOrNull(symbols: List): Expression? +public interface DifferentiableExpression> : Expression { + /** + * Differentiates this expression by ordered collection of [symbols]. + */ + public fun derivativeOrNull(symbols: List): R? } -public fun DifferentiableExpression.derivative(symbols: List): Expression = +public fun > DifferentiableExpression.derivative(symbols: List): R = derivativeOrNull(symbols) ?: error("Derivative by symbols $symbols not provided") -public fun DifferentiableExpression.derivative(vararg symbols: Symbol): Expression = +public fun > DifferentiableExpression.derivative(vararg symbols: Symbol): R = derivative(symbols.toList()) -public fun DifferentiableExpression.derivative(name: String): Expression = +public fun > DifferentiableExpression.derivative(name: String): R = derivative(StringSymbol(name)) /** * A [DifferentiableExpression] that defines only first derivatives */ -public abstract class FirstDerivativeExpression : DifferentiableExpression { +public abstract class FirstDerivativeExpression> : DifferentiableExpression { + /** + * Returns first derivative of this expression by given [symbol]. + */ + public abstract fun derivativeOrNull(symbol: Symbol): R? - public abstract fun derivativeOrNull(symbol: Symbol): Expression? - - public override fun derivativeOrNull(symbols: List): Expression? { + public final override fun derivativeOrNull(symbols: List): R? { val dSymbol = symbols.firstOrNull() ?: return null return derivativeOrNull(dSymbol) } @@ -32,6 +40,6 @@ public abstract class FirstDerivativeExpression : DifferentiableExpression /** * A factory that converts an expression in autodiff variables to a [DifferentiableExpression] */ -public interface AutoDiffProcessor> { - public fun process(function: A.() -> I): DifferentiableExpression -} \ No newline at end of file +public fun interface AutoDiffProcessor, R : Expression> { + public fun process(function: A.() -> I): DifferentiableExpression +} diff --git a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/Expression.kt b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/Expression.kt index 568de255e..98940e767 100644 --- a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/Expression.kt +++ b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/Expression.kt @@ -22,7 +22,9 @@ public inline class StringSymbol(override val identity: String) : Symbol { } /** - * An elementary function that could be invoked on a map of arguments + * An elementary function that could be invoked on a map of arguments. + * + * @param T the type this expression takes as argument and returns. */ public fun interface Expression { /** diff --git a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/SimpleAutoDiff.kt b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/SimpleAutoDiff.kt index 5a9642690..e8a894d23 100644 --- a/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/SimpleAutoDiff.kt +++ b/kmath-core/src/commonMain/kotlin/kscience/kmath/expressions/SimpleAutoDiff.kt @@ -68,7 +68,7 @@ public fun > F.simpleAutoDiff( ): DerivationResult { contract { callsInPlace(body, InvocationKind.EXACTLY_ONCE) } - return SimpleAutoDiffField(this, bindings).derivate(body) + return SimpleAutoDiffField(this, bindings).differentiate(body) } public fun > F.simpleAutoDiff( @@ -83,12 +83,21 @@ public open class SimpleAutoDiffField>( public val context: F, bindings: Map, ) : Field>, ExpressionAlgebra> { + public override val zero: AutoDiffValue + get() = const(context.zero) + + public override val one: AutoDiffValue + get() = const(context.one) // this stack contains pairs of blocks and values to apply them to private var stack: Array = arrayOfNulls(8) private var sp: Int = 0 private val derivatives: MutableMap, T> = hashMapOf() + private val bindings: Map> = bindings.entries.associate { + it.key.identity to AutoDiffVariableWithDerivative(it.key.identity, it.value, context.zero) + } + /** * Differentiable variable with value and derivative of differentiation ([simpleAutoDiff]) result * with respect to this variable. @@ -106,11 +115,7 @@ public open class SimpleAutoDiffField>( override fun hashCode(): Int = identity.hashCode() } - private val bindings: Map> = bindings.entries.associate { - it.key.identity to AutoDiffVariableWithDerivative(it.key.identity, it.value, context.zero) - } - - override fun bindOrNull(symbol: Symbol): AutoDiffValue? = bindings[symbol.identity] + public override fun bindOrNull(symbol: Symbol): AutoDiffValue? = bindings[symbol.identity] private fun getDerivative(variable: AutoDiffValue): T = (variable as? AutoDiffVariableWithDerivative)?.d ?: derivatives[variable] ?: context.zero @@ -119,7 +124,6 @@ public open class SimpleAutoDiffField>( if (variable is AutoDiffVariableWithDerivative) variable.d = value else derivatives[variable] = value } - @Suppress("UNCHECKED_CAST") private fun runBackwardPass() { while (sp > 0) { @@ -129,9 +133,6 @@ public open class SimpleAutoDiffField>( } } - override val zero: AutoDiffValue get() = const(context.zero) - override val one: AutoDiffValue get() = const(context.one) - override fun const(value: T): AutoDiffValue = AutoDiffValue(value) /** @@ -165,7 +166,7 @@ public open class SimpleAutoDiffField>( } - internal fun derivate(function: SimpleAutoDiffField.() -> AutoDiffValue): DerivationResult { + internal fun differentiate(function: SimpleAutoDiffField.() -> AutoDiffValue): DerivationResult { val result = function() result.d = context.one // computing derivative w.r.t result runBackwardPass() @@ -174,41 +175,41 @@ public open class SimpleAutoDiffField>( // Overloads for Double constants - override operator fun Number.plus(b: AutoDiffValue): AutoDiffValue = + public override operator fun Number.plus(b: AutoDiffValue): AutoDiffValue = derive(const { this@plus.toDouble() * one + b.value }) { z -> b.d += z.d } - override operator fun AutoDiffValue.plus(b: Number): AutoDiffValue = b.plus(this) + public override operator fun AutoDiffValue.plus(b: Number): AutoDiffValue = b.plus(this) - override operator fun Number.minus(b: AutoDiffValue): AutoDiffValue = + public override operator fun Number.minus(b: AutoDiffValue): AutoDiffValue = derive(const { this@minus.toDouble() * one - b.value }) { z -> b.d -= z.d } - override operator fun AutoDiffValue.minus(b: Number): AutoDiffValue = + public override operator fun AutoDiffValue.minus(b: Number): AutoDiffValue = derive(const { this@minus.value - one * b.toDouble() }) { z -> this@minus.d += z.d } // Basic math (+, -, *, /) - override fun add(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = + public override fun add(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = derive(const { a.value + b.value }) { z -> a.d += z.d b.d += z.d } - override fun multiply(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = + public override fun multiply(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = derive(const { a.value * b.value }) { z -> a.d += z.d * b.value b.d += z.d * a.value } - override fun divide(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = + public override fun divide(a: AutoDiffValue, b: AutoDiffValue): AutoDiffValue = derive(const { a.value / b.value }) { z -> a.d += z.d / b.value b.d -= z.d * a.value / (b.value * b.value) } - override fun multiply(a: AutoDiffValue, k: Number): AutoDiffValue = + public override fun multiply(a: AutoDiffValue, k: Number): AutoDiffValue = derive(const { k.toDouble() * a.value }) { z -> a.d += z.d * k.toDouble() } @@ -220,15 +221,15 @@ public open class SimpleAutoDiffField>( public class SimpleAutoDiffExpression>( public val field: F, public val function: SimpleAutoDiffField.() -> AutoDiffValue, -) : FirstDerivativeExpression() { +) : FirstDerivativeExpression>() { public override operator fun invoke(arguments: Map): T { //val bindings = arguments.entries.map { it.key.bind(it.value) } return SimpleAutoDiffField(field, arguments).function().value } - override fun derivativeOrNull(symbol: Symbol): Expression = Expression { arguments -> + public override fun derivativeOrNull(symbol: Symbol): Expression = Expression { arguments -> //val bindings = arguments.entries.map { it.key.bind(it.value) } - val derivationResult = SimpleAutoDiffField(field, arguments).derivate(function) + val derivationResult = SimpleAutoDiffField(field, arguments).differentiate(function) derivationResult.derivative(symbol) } } @@ -236,13 +237,10 @@ public class SimpleAutoDiffExpression>( /** * Generate [AutoDiffProcessor] for [SimpleAutoDiffExpression] */ -public fun > simpleAutoDiff(field: F): AutoDiffProcessor, SimpleAutoDiffField> { - return object : AutoDiffProcessor, SimpleAutoDiffField> { - override fun process(function: SimpleAutoDiffField.() -> AutoDiffValue): DifferentiableExpression { - return SimpleAutoDiffExpression(field, function) - } +public fun > simpleAutoDiff(field: F): AutoDiffProcessor, SimpleAutoDiffField, Expression> = + AutoDiffProcessor { function -> + SimpleAutoDiffExpression(field, function) } -} // Extensions for differentiation of various basic mathematical functions @@ -392,4 +390,4 @@ public class SimpleAutoDiffExtendedField>( public override fun atanh(arg: AutoDiffValue): AutoDiffValue = (this as SimpleAutoDiffField).atanh(arg) -} \ No newline at end of file +} diff --git a/kmath-kotlingrad/src/main/kotlin/kscience/kmath/kotlingrad/DifferentiableMstExpression.kt b/kmath-kotlingrad/src/main/kotlin/kscience/kmath/kotlingrad/DifferentiableMstExpression.kt index ddfc7cccb..dd5e46f90 100644 --- a/kmath-kotlingrad/src/main/kotlin/kscience/kmath/kotlingrad/DifferentiableMstExpression.kt +++ b/kmath-kotlingrad/src/main/kotlin/kscience/kmath/kotlingrad/DifferentiableMstExpression.kt @@ -5,7 +5,6 @@ import kscience.kmath.ast.MST import kscience.kmath.ast.MstAlgebra import kscience.kmath.ast.MstExpression import kscience.kmath.expressions.DifferentiableExpression -import kscience.kmath.expressions.StringSymbol import kscience.kmath.expressions.Symbol import kscience.kmath.operations.NumericAlgebra @@ -20,7 +19,7 @@ import kscience.kmath.operations.NumericAlgebra * @property expr the underlying [MstExpression]. */ public inline class DifferentiableMstExpression(public val expr: MstExpression) : - DifferentiableExpression where A : NumericAlgebra, T : Number { + DifferentiableExpression> where A : NumericAlgebra, T : Number { public constructor(algebra: A, mst: MST) : this(MstExpression(algebra, mst)) /** @@ -47,15 +46,6 @@ public inline class DifferentiableMstExpression(public val expr: MstExpres ) } -public fun > DifferentiableMstExpression.derivative(symbols: List): MstExpression = - derivativeOrNull(symbols) - -public fun > DifferentiableMstExpression.derivative(vararg symbols: Symbol): MstExpression = - derivative(symbols.toList()) - -public fun > DifferentiableMstExpression.derivative(name: String): MstExpression = - derivative(StringSymbol(name)) - /** * Wraps this [MstExpression] into [DifferentiableMstExpression]. */