1 files changed, 148 insertions, 0 deletions
diff --git a/src/main/java/org/apache/commons/math3/ode/FieldExpandableODE.java b/src/main/java/org/apache/commons/math3/ode/FieldExpandableODE.java
new file mode 100644
index 0000000..a9f079e
--- /dev/null
+++ b/src/main/java/org/apache/commons/math3/ode/FieldExpandableODE.java
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF 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
+ *
+ *      http://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.apache.commons.math3.ode;
+
+import org.apache.commons.math3.RealFieldElement;
+import org.apache.commons.math3.exception.DimensionMismatchException;
+import org.apache.commons.math3.exception.MaxCountExceededException;
+import org.apache.commons.math3.util.MathArrays;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * This class represents a combined set of first order differential equations, with at least a
+ * primary set of equations expandable by some sets of secondary equations.
+ *
+ * <p>One typical use case is the computation of the Jacobian matrix for some ODE. In this case, the
+ * primary set of equations corresponds to the raw ODE, and we add to this set another bunch of
+ * secondary equations which represent the Jacobian matrix of the primary set.
+ *
+ * <p>We want the integrator to use <em>only</em> the primary set to estimate the errors and hence
+ * the step sizes. It should <em>not</em> use the secondary equations in this computation. The
+ * {@link FirstOrderFieldIntegrator integrator} will be able to know where the primary set ends and
+ * so where the secondary sets begin.
+ *
+ * @see FirstOrderFieldDifferentialEquations
+ * @see FieldSecondaryEquations
+ * @param <T> the type of the field elements
+ * @since 3.6
+ */
+public class FieldExpandableODE<T extends RealFieldElement<T>> {
+
+    /** Primary differential equation. */
+    private final FirstOrderFieldDifferentialEquations<T> primary;
+
+    /** Components of the expandable ODE. */
+    private List<FieldSecondaryEquations<T>> components;
+
+    /** Mapper for all equations. */
+    private FieldEquationsMapper<T> mapper;
+
+    /**
+     * Build an expandable set from its primary ODE set.
+     *
+     * @param primary the primary set of differential equations to be integrated.
+     */
+    public FieldExpandableODE(final FirstOrderFieldDifferentialEquations<T> primary) {
+        this.primary = primary;
+        this.components = new ArrayList<FieldSecondaryEquations<T>>();
+        this.mapper = new FieldEquationsMapper<T>(null, primary.getDimension());
+    }
+
+    /**
+     * Get the mapper for the set of equations.
+     *
+     * @return mapper for the set of equations
+     */
+    public FieldEquationsMapper<T> getMapper() {
+        return mapper;
+    }
+
+    /**
+     * Add a set of secondary equations to be integrated along with the primary set.
+     *
+     * @param secondary secondary equations set
+     * @return index of the secondary equation in the expanded state, to be used as the parameter to
+     *     {@link FieldODEState#getSecondaryState(int)} and {@link
+     *     FieldODEStateAndDerivative#getSecondaryDerivative(int)} (beware index 0 corresponds to
+     *     main state, additional states start at 1)
+     */
+    public int addSecondaryEquations(final FieldSecondaryEquations<T> secondary) {
+
+        components.add(secondary);
+        mapper = new FieldEquationsMapper<T>(mapper, secondary.getDimension());
+
+        return components.size();
+    }
+
+    /**
+     * Initialize equations at the start of an ODE integration.
+     *
+     * @param t0 value of the independent <I>time</I> variable at integration start
+     * @param y0 array containing the value of the state vector at integration start
+     * @param finalTime target time for the integration
+     * @exception MaxCountExceededException if the number of functions evaluations is exceeded
+     * @exception DimensionMismatchException if arrays dimensions do not match equations settings
+     */
+    public void init(final T t0, final T[] y0, final T finalTime) {
+
+        // initialize primary equations
+        int index = 0;
+        final T[] primary0 = mapper.extractEquationData(index, y0);
+        primary.init(t0, primary0, finalTime);
+
+        // initialize secondary equations
+        while (++index < mapper.getNumberOfEquations()) {
+            final T[] secondary0 = mapper.extractEquationData(index, y0);
+            components.get(index - 1).init(t0, primary0, secondary0, finalTime);
+        }
+    }
+
+    /**
+     * Get the current time derivative of the complete state vector.
+     *
+     * @param t current value of the independent <I>time</I> variable
+     * @param y array containing the current value of the complete state vector
+     * @return time derivative of the complete state vector
+     * @exception MaxCountExceededException if the number of functions evaluations is exceeded
+     * @exception DimensionMismatchException if arrays dimensions do not match equations settings
+     */
+    public T[] computeDerivatives(final T t, final T[] y)
+            throws MaxCountExceededException, DimensionMismatchException {
+
+        final T[] yDot = MathArrays.buildArray(t.getField(), mapper.getTotalDimension());
+
+        // compute derivatives of the primary equations
+        int index = 0;
+        final T[] primaryState = mapper.extractEquationData(index, y);
+        final T[] primaryStateDot = primary.computeDerivatives(t, primaryState);
+        mapper.insertEquationData(index, primaryStateDot, yDot);
+
+        // Add contribution for secondary equations
+        while (++index < mapper.getNumberOfEquations()) {
+            final T[] componentState = mapper.extractEquationData(index, y);
+            final T[] componentStateDot =
+                    components
+                            .get(index - 1)
+                            .computeDerivatives(t, primaryState, primaryStateDot, componentState);
+            mapper.insertEquationData(index, componentStateDot, yDot);
+        }
+
+        return yDot;
+    }
+}