VariationalEquation.java

/*
 * Licensed to the Hipparchus project under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The Hipparchus project licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.hipparchus.ode;

import java.lang.reflect.Array;

import org.hipparchus.exception.LocalizedCoreFormats;
import org.hipparchus.exception.MathIllegalArgumentException;
import org.hipparchus.exception.MathIllegalStateException;

/**
 * This class defines a set of {@link SecondaryODE secondary equations} to
 * compute the global Jacobian matrices with respect to the initial state
 * vector and, if any, to some parameters of the primary ODE set.
 * <p>
 * The primary set of ODE for which Jaobian matrices are requested may be:
 * </p>
 * <ul>
 * <li>a full-fledged {@link ODEJacobiansProvider} that computes by itself
 * both the ODE and its local partial derivatives,</li>
 * <li>a simple {@link OrdinaryDifferentialEquation} which must therefore
 * be completed with a finite differences configuration to compute local
 * partial derivatives (so-called internal differentiation).</li>
 * </ul>
 * <p>
 * As the variational equation automatically inserts {@link
 * ExpandableODE#addSecondaryEquations(SecondaryODE) secondary differential
 * equations}, in the {@link ExpandableODE expandable ODE}, data for
 * initial state must also be inserted before integration and matrices
 * result must be extracted after integration. This implies a precise
 * scheduling of the calls to the various methods of this class. The
 * proper scheduling is the following one:
 * </p>
 * <pre>
 *   // set up equations
 *   ODEJacobiansProvider jode       = new MyODE(...);
 *   ExpandableODE        expandable = new Expandable(jode);
 *   VariationalEquation  ve         = new VariationalEquation(expandable, jode);
 *
 *   // set up initial state
 *   ODEState initWithoutDerivatives = new ODEState(t0, y0);
 *   ve.setInitialMainStateJacobian(dYdY0); // only needed if the default identity matrix is not suitable
 *   ve.setInitialParameterJacobian(name, dYdP); // only needed if the default zero matrix is not suitable
 *   ODEState initWithDerivatives = ve.setUpInitialState(initWithoutDerivatives);
 *
 *   // perform integration on the expanded equations with the expanded initial state
 *   ODEStateAndDerivative finalState = integrator.integrate(expandable, initWithDerivatives, finalT);
 *
 *   // extract Jacobian matrices
 *   dYdY0 = ve.extractMainSetJacobian(finalState);
 *   dYdP  = ve.extractParameterJacobian(finalState, name);
 * </pre>
 * <p>
 * The most important part is to not forget to call {@link #setUpInitialState(ODEState)} to add
 * the secondary state with the initial matrices to the {@link ODEState} used in the
 * {@link ODEIntegrator#integrate(ExpandableODE, ODEState, double) integrate} method.
 * Forgetting to do this and passing only a {@link ODEState} without the secondary state
 * set up will trigger an error as the state vector will not have the correct dimension.
 * </p>
 *
 * @see ExpandableODE
 * @see ODEJacobiansProvider
 * @see OrdinaryDifferentialEquation
 * @see NamedParameterJacobianProvider
 * @see ParametersController
 *
 */
public class VariationalEquation {

    /** ODE with Jacobian computation skill. */
    private final ODEJacobiansProvider jode;

    /** Expandable first order differential equation. */
    private final ExpandableODE expandable;

    /** Index of the instance in the expandable set. */
    private final int index;

    /** State and parameters Jacobian matrices in a row. */
    private double[] matricesData;

    /** Build variational equation using finite differences for local
     * partial derivatives.
     * @param expandable expandable set into which variational equations should be registered
     * @param ode base ordinary differential equation for which Jacobians
     * matrices are requested
     * @param hY step used for finite difference computation with respect to state vector
     * @param controller controller to change parameters
     * @param paramsAndSteps parameters and steps to compute the Jacobians df/dp
     * @exception MismatchedEquations if the primary set of the expandable set does
     * not match the {@code ode}
     */
    public VariationalEquation(final ExpandableODE expandable,
                               final OrdinaryDifferentialEquation ode, final double[] hY,
                               final ParametersController controller,
                               final ParameterConfiguration ... paramsAndSteps)
        throws MismatchedEquations {
        this(expandable, new ParameterJacobianWrapper(ode, hY, controller, paramsAndSteps));
    }

    /** Build variational equation using analytical local partial derivatives.
     * <p>
     * Parameters must belong to the supported ones given by {@link
     * Parameterizable#getParametersNames()}, so the primary set of differential
     * equations must be {@link Parameterizable}.
     * </p>
     * <p>Note that each selection clears the previous selected parameters.</p>
     *
     * @param expandable expandable set into which variational equations should be registered
     * @param jode the primary first order differential equations set to extend
     * @exception MismatchedEquations if the primary set of the expandable set does
     * not match the {@code ode}
     */
    public VariationalEquation(final ExpandableODE expandable,
                               final ODEJacobiansProvider jode)
        throws MismatchedEquations {

        // safety checks
        final OrdinaryDifferentialEquation ode;
        if (jode instanceof ParameterJacobianWrapper) {
            ode = ((ParameterJacobianWrapper) jode).getODE();
        } else {
            ode = jode;
        }
        if (expandable.getPrimary() != ode) {
            throw new MismatchedEquations();
        }

        this.jode       = jode;
        this.expandable = expandable;
        this.index      = expandable.addSecondaryEquations(new JacobiansSecondaryODE());

        // set the default initial state Jacobian to the identity
        // and the default initial parameters Jacobian to the null matrix
        matricesData = new double[(jode.getDimension() + jode.getParametersNames().size()) * jode.getDimension()];
        for (int i = 0; i < jode.getDimension(); ++i) {
            matricesData[i * (jode.getDimension() + 1)] = 1.0;
        }

    }

    /** Set the initial value of the Jacobian matrix with respect to state.
     * <p>
     * If this method is not called, the initial value of the Jacobian
     * matrix with respect to state is set to identity.
     * </p>
     * <p>
     * This method must be called <em>before {@link #setUpInitialState(ODEState)}</em>
     * </p>
     * @param dYdY0 initial Jacobian matrix w.r.t. state
     * @exception MathIllegalArgumentException if matrix dimensions are incorrect
     */
    public void setInitialMainStateJacobian(final double[][] dYdY0)
        throws MathIllegalArgumentException {

        // Check dimensions
        checkDimension(jode.getDimension(), dYdY0);
        checkDimension(jode.getDimension(), dYdY0[0]);

        // store the matrix in row major order as a single dimension array
        int i = 0;
        for (final double[] row : dYdY0) {
            System.arraycopy(row, 0, matricesData, i, jode.getDimension());
            i += jode.getDimension();
        }

    }

    /** Set the initial value of a column of the Jacobian matrix with respect to one parameter.
     * <p>
     * If this method is not called for some parameter, the initial value of
     * the column of the Jacobian matrix with respect to this parameter is set to zero.
     * </p>
     * <p>
     * This method must be called <em>before {@link #setUpInitialState(ODEState)}</em>
     * </p>
     * @param pName parameter name
     * @param dYdP initial Jacobian column vector with respect to the parameter
     * @exception MathIllegalArgumentException if a parameter is not supported
     * @throws MathIllegalArgumentException if the column vector does not match state dimension
     */
    public void setInitialParameterJacobian(final String pName, final double[] dYdP)
        throws MathIllegalArgumentException {

        // Check dimensions
        checkDimension(jode.getDimension(), dYdP);

        // store the column in a global single dimension array
        int i = jode.getDimension() * jode.getDimension();
        for (final String knownParameter : jode.getParametersNames()) {
            if (pName.equals(knownParameter)) {
                System.arraycopy(dYdP, 0, matricesData, i, jode.getDimension());
                return;
            }
            i += jode.getDimension();
        }

        throw new MathIllegalArgumentException(LocalizedODEFormats.UNKNOWN_PARAMETER, pName);

    }

    /** Set up initial state.
     * <p>
     * This method inserts the initial Jacobian matrices data into
     * an {@link ODEState ODE state} by overriding the additional
     * state components corresponding to the instance. It must be
     * called prior to integrate the equations.
     * </p>
     * <p>
     * This method must be called <em>after</em>
     * {@link #setInitialMainStateJacobian(double[][])} and
     * {@link #setInitialParameterJacobian(String, double[])}.
     * </p>
     * @param initialState initial state, without the initial Jacobians
     * matrices
     * @return a new instance of initial state, with the initial Jacobians
     * matrices properly initialized
     */
    public ODEState setUpInitialState(final ODEState initialState) { // NOPMD - PMD false positive

        // insert the matrices data into secondary states
        final double[][] secondary = new double[expandable.getMapper().getNumberOfEquations() - 1][];
        for (int i = 0; i < initialState.getNumberOfSecondaryStates(); ++i) {
            if (i + 1 != index) {
                secondary[i] = initialState.getSecondaryState(i + 1);
            }
        }
        secondary[index - 1] = matricesData;

        // create an updated initial state
        return new ODEState(initialState.getTime(), initialState.getPrimaryState(), secondary);

    }

    /** Extract the Jacobian matrix with respect to state.
     * @param state state from which to extract Jacobian matrix
     * @return Jacobian matrix dY/dY0 with respect to state.
     */
    public double[][] extractMainSetJacobian(final ODEState state) {

        // get current state for this set of equations from the expandable fode
        final double[] p = state.getSecondaryState(index);

        final double[][] dYdY0 = new double[jode.getDimension()][jode.getDimension()];
        int j = 0;
        for (int i = 0; i < jode.getDimension(); i++) {
            System.arraycopy(p, j, dYdY0[i], 0, jode.getDimension());
            j += jode.getDimension();
        }

        return dYdY0;

    }

    /** Extract the Jacobian matrix with respect to one parameter.
     * @param state state from which to extract Jacobian matrix
     * @param pName name of the parameter for the computed Jacobian matrix
     * @return Jacobian matrix dY/dP with respect to the named parameter
     */
    public double[] extractParameterJacobian(final ODEState state, final String pName) {

        // get current state for this set of equations from the expandable fode
        final double[] p = state.getSecondaryState(index);

        final double[] dYdP = new double[jode.getDimension()];
        int i = jode.getDimension() * jode.getDimension();
        for (final String knownParameter : jode.getParametersNames()) {
            if (pName.equals(knownParameter)) {
                System.arraycopy(p, i, dYdP, 0, jode.getDimension());
                break;
            }
            i += jode.getDimension();
        }

        return dYdP;

    }

    /** Check array dimensions.
     * @param expected expected dimension
     * @param array (may be null if expected is 0)
     * @throws MathIllegalArgumentException if the array dimension does not match the expected one
     */
    private void checkDimension(final int expected, final Object array)
        throws MathIllegalArgumentException {
        int arrayDimension = (array == null) ? 0 : Array.getLength(array);
        if (arrayDimension != expected) {
            throw new MathIllegalArgumentException(LocalizedCoreFormats.DIMENSIONS_MISMATCH,
                                                   arrayDimension, expected);
        }
    }

    /** Local implementation of secondary equations. */
    private class JacobiansSecondaryODE implements SecondaryODE {

        /** {@inheritDoc} */
        @Override
        public int getDimension() {
            return jode.getDimension() * (jode.getDimension() + jode.getParametersNames().size());
        }

        /** {@inheritDoc} */
        @Override
        public double[] computeDerivatives(final double t, final double[] y, final double[] yDot,
                                           final double[] z)
            throws MathIllegalArgumentException, MathIllegalStateException {

            final double[] zDot = new double[z.length];

            // variational equations:
            // from d[dy/dt]/dy0 and d[dy/dt]/dp to d[dy/dy0]/dt and d[dy/dp]/dt

            // compute Jacobian matrix with respect to primary state
            double[][] dFdY = jode.computeMainStateJacobian(t, y, yDot);

            // Dispatch Jacobian matrix in the compound secondary state vector
            for (int i = 0; i < jode.getDimension(); ++i) {
                final double[] dFdYi = dFdY[i];
                for (int j = 0; j < jode.getDimension(); ++j) {
                    double s = 0;
                    final int startIndex = j;
                    int zIndex = startIndex;
                    for (int l = 0; l < jode.getDimension(); ++l) {
                        s += dFdYi[l] * z[zIndex];
                        zIndex += jode.getDimension();
                    }
                    zDot[startIndex + i * jode.getDimension()] = s;
                }
            }

            // compute Jacobian matrices with respect to parameters
            int startIndex = jode.getDimension() * jode.getDimension();
            for (final String name : jode.getParametersNames()) {
                final double[] dFdP = jode.computeParameterJacobian(t, y, yDot, name);
                for (int i = 0; i < jode.getDimension(); ++i) {
                    final double[] dFdYi = dFdY[i];
                    int zIndex = startIndex;
                    double s = dFdP[i];
                    for (int l = 0; l < jode.getDimension(); ++l) {
                        s += dFdYi[l] * z[zIndex];
                        zIndex++;
                    }
                    zDot[startIndex + i] = s;
                }
                startIndex += jode.getDimension();
            }

            return zDot;

        }
    }

    /**
     * Special exception for equations mismatch.
     */
    public static class MismatchedEquations extends MathIllegalArgumentException {

        /** Serializable UID. */
        private static final long serialVersionUID = 20120902L;

        /** Simple constructor. */
        public MismatchedEquations() {
            super(LocalizedODEFormats.UNMATCHED_ODE_IN_EXPANDED_SET);
        }

    }

}