Add efficient resource reduction

Adds reduce methods to a number of Resource monads, which do a "flat" resource reduction. That is, they use ThinResourceMonad.reduce to produce a single resource node, which collects all dependencies and reduces the dynamics, rather than the original Resources.reduce method, which creates a number of intermediate resources linear in the number of operands. When doing large roll-ups, like a power model, those intermediate nodes each contribute a small performance overhead which can eventually add up. Providing a standard way to do reduction efficiently, and using that to do common operations like taking a sum or maximum, should help.
NASA-AMMOS · Sep 3, 2024 · 3135d23 · 3135d23
1 parent d353474
commit 3135d23
Show file tree

Hide file tree

Showing 8 changed files with 190 additions and 36 deletions.
diff --git a/contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/Resources.java b/contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/Resources.java
@@ -304,10 +304,37 @@ public static <D> Resource<D> eraseExpiry(Resource<D> p) {
     return result;
   }
 
+  /**
+   * Reduce a collection of resources.
+   * <p>
+   *     Note on efficiency: For n operands, this method will create (n-2) intermediate resources,
+   *     plus the terminal result resource.
+   *     This overhead is often unnecessary. Consider using the methods linked below
+   *     for alternative implementations of reduce that may be more efficient.
+   *     There may also be efficient implementations of reduce on the relevant resource monad, like
+   *     {@link gov.nasa.jpl.aerie.contrib.streamline.modeling.discrete.monads.DiscreteResourceMonad#reduce}.
+   * </p>
+   *
+   * @see gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   */
   public static <D> Resource<D> reduce(Stream<? extends Resource<D>> operands, Resource<D> identity, BiFunction<Resource<D>, Resource<D>, Resource<D>> operation, String operationName) {
     return reduce(operands.toList(), identity, operation, operationName);
   }
 
+  /**
+   * Reduce a collection of resources.
+   * <p>
+   *     Note on efficiency: For n operands, this method will create O(n) intermediate resources.
+   *     This overhead is often unnecessary. Consider using the methods linked below
+   *     for alternative implementations of reduce that may be more efficient.
+   *     There may also be efficient implementations of reduce on the relevant resource monad, like
+   *     {@link gov.nasa.jpl.aerie.contrib.streamline.modeling.discrete.monads.DiscreteResourceMonad#reduce}.
+   * </p>
+   *
+   * @see gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   */
   public static <D> Resource<D> reduce(Collection<? extends Resource<D>> operands, Resource<D> identity, BiFunction<Resource<D>, Resource<D>, Resource<D>> operation, String operationName) {
     var result = operands.stream().reduce(identity, operation, operation::apply);
     name(result, operationName + argsFormat(operands), operands.toArray());

diff --git a/contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/monads/ResourceMonad.java b/contrib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/monads/ResourceMonad.java
@@ -8,10 +8,13 @@
 import gov.nasa.jpl.aerie.contrib.streamline.utils.*;
 import org.apache.commons.lang3.function.TriFunction;
 
+import java.util.Collection;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Dependencies.addDependency;
+import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Naming.argsFormat;
+import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Naming.name;
 import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Profiling.profile;
 import static gov.nasa.jpl.aerie.contrib.streamline.utils.FunctionalUtils.curry;
 
@@ -72,6 +75,63 @@ public static <A> Resource<A> join(Resource<Resource<A>> a) {
     return result;
   }
 
+  /**
+   * Efficient reduce for resources, lifting an operator that can reduce the dynamics.
+   * <p>
+   *     This is logically equivalent to, but more efficient than,
+   *     <pre>operands.stream().reduce(pure(identity), map(operator), map(operator))</pre>
+   *     That is, it's logically equivalent to lifting the operator to an operator on resources,
+   *     then reducing the resources with the lifted operator.
+   *     However, that would produce a large number of unnecessary intermediate resources.
+   *     This function creates a "flat" reduction, with no intermediate nodes.
+   * </p>
+   *
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction, String)
+   */
+  public static <A> Resource<A> reduce(Collection<? extends Resource<A>> operands, A identity, BiFunction<A, A, A> f) {
+    return reduce(operands, DynamicsMonad.pure(identity), DynamicsMonad.map(f));
+  }
+
+  /**
+   * Like {@link ResourceMonad#reduce(Collection, Object, BiFunction)}, but also names the result
+   * like "operationName(operand1, operand2, ...)".
+   *
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction, String)
+   */
+  public static <A> Resource<A> reduce(Collection<? extends Resource<A>> operands, A identity, BiFunction<A, A, A> f, String operationName) {
+    return reduce(operands, DynamicsMonad.pure(identity), DynamicsMonad.map(f), operationName);
+  }
+
+  /**
+   * Like {@link ResourceMonad#reduce(Collection, Object, BiFunction)},
+   * but operator acts on fully wrapped dynamics instead of plain dynamics.
+   *
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction, String)
+   */
+  public static <A> Resource<A> reduce(Collection<? extends Resource<A>> operands, ErrorCatching<Expiring<A>> identity, BiFunction<ErrorCatching<Expiring<A>>, ErrorCatching<Expiring<A>>, ErrorCatching<Expiring<A>>> f) {
+    Resource<A> result = ThinResourceMonad.reduce(operands, identity, f)::getDynamics;
+    operands.forEach(op -> addDependency(result, op));
+    return result;
+  }
+
+  /**
+   * Like {@link ResourceMonad#reduce(Collection, Object, BiFunction, String)},
+   * but operator acts on fully wrapped dynamics instead of plain dynamics.
+   *
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   * @see ResourceMonad#reduce(Collection, ErrorCatching, BiFunction)
+   */
+  public static <A> Resource<A> reduce(Collection<? extends Resource<A>> operands, ErrorCatching<Expiring<A>> identity, BiFunction<ErrorCatching<Expiring<A>>, ErrorCatching<Expiring<A>>, ErrorCatching<Expiring<A>>> f, String operationName) {
+    return name(reduce(operands, identity, f), operationName + argsFormat(operands), operands.toArray());
+  }
+
   // Not strictly part of this monad, but commonly used to "fill the gap" when deriving resources with partial bindings
   public static <A> Resource<A> pure(Expiring<A> a) {
     return ThinResourceMonad.pure(ErrorCatchingMonad.pure(a))::getDynamics;

diff --git a/...ib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/monads/ThinResourceMonad.java b/...ib/src/main/java/gov/nasa/jpl/aerie/contrib/streamline/core/monads/ThinResourceMonad.java
@@ -4,6 +4,7 @@
 import gov.nasa.jpl.aerie.contrib.streamline.utils.*;
 import org.apache.commons.lang3.function.TriFunction;
 
+import java.util.Collection;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -24,6 +25,18 @@ public static <A> ThinResource<A> join(ThinResource<ThinResource<A>> a) {
     return () -> a.getDynamics().getDynamics();
   }
 
+  /**
+   * Efficiently reduce {@link ThinResource}s.
+   * <p>
+   *     Note: This is a lower-level function that doesn't record dependencies,
+   *     meant to be used in a select few cases.
+   *     Most users would prefer {@link ResourceMonad#reduce}.
+   * </p>
+   */
+  public static <A> ThinResource<A> reduce(Collection<? extends ThinResource<A>> operands, A identity, BiFunction<A, A, A> f) {
+    return () -> operands.stream().map(ThinResource::getDynamics).reduce(identity, f, f::apply);
+  }
+
   // GENERATED CODE START
   // Supplemental methods generated by generate_monad_methods.py on 2023-12-06.
 

diff --git a/...ava/gov/nasa/jpl/aerie/contrib/streamline/modeling/black_box/DifferentiableResources.java b/...ava/gov/nasa/jpl/aerie/contrib/streamline/modeling/black_box/DifferentiableResources.java
@@ -5,10 +5,10 @@
 import gov.nasa.jpl.aerie.contrib.streamline.modeling.polynomial.Polynomial;
 import gov.nasa.jpl.aerie.merlin.protocol.types.Duration;
 
-import java.util.stream.Stream;
+import java.util.Collection;
 
-import static gov.nasa.jpl.aerie.contrib.streamline.core.Resources.reduce;
 import static gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad.map;
+import static gov.nasa.jpl.aerie.contrib.streamline.core.monads.ResourceMonad.reduce;
 import static gov.nasa.jpl.aerie.contrib.streamline.debugging.Naming.name;
 import static gov.nasa.jpl.aerie.contrib.streamline.modeling.black_box.Differentiable.differentiable;
 import static java.util.Arrays.stream;
@@ -49,11 +49,11 @@ public static Resource<Differentiable> asDifferentiable(Resource<Polynomial> pol
 
     @SafeVarargs
     public static Resource<Differentiable> add(Resource<Differentiable>... summands) {
-        return sum(stream(summands));
+        return sum(stream(summands).toList());
     }
 
-    public static Resource<Differentiable> sum(Stream<Resource<Differentiable>> summands) {
-        return reduce(summands, constant(0), map(Differentiable::add), "Sum");
+    public static Resource<Differentiable> sum(Collection<Resource<Differentiable>> summands) {
+        return reduce(summands, Differentiable.constant(0), Differentiable::add, "Sum");
     }
 
     public static Resource<Differentiable> subtract(Resource<Differentiable> left, Resource<Differentiable> right) {
@@ -64,11 +64,11 @@ public static Resource<Differentiable> subtract(Resource<Differentiable> left, R
 
     @SafeVarargs
     public static Resource<Differentiable> multiply(Resource<Differentiable>... factors) {
-        return product(stream(factors));
+        return product(stream(factors).toList());
     }
 
-    public static Resource<Differentiable> product(Stream<Resource<Differentiable>> factors) {
-        return reduce(factors, constant(0), map(Differentiable::multiply), "Product");
+    public static Resource<Differentiable> product(Collection<Resource<Differentiable>> factors) {
+        return reduce(factors, Differentiable.constant(0), Differentiable::multiply, "Product");
     }
 
     public static Resource<Differentiable> divide(Resource<Differentiable> left, Resource<Differentiable> right) {

diff --git a/...l/aerie/contrib/streamline/modeling/black_box/monads/UnstructuredResourceApplicative.java b/...l/aerie/contrib/streamline/modeling/black_box/monads/UnstructuredResourceApplicative.java
@@ -6,6 +6,7 @@
 import gov.nasa.jpl.aerie.contrib.streamline.utils.*;
 import org.apache.commons.lang3.function.TriFunction;
 
+import java.util.Collection;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -37,6 +38,26 @@ public static <A, B> Resource<Unstructured<B>> apply(Resource<Unstructured<A>> a
     //     return ResourceMonad.map(ResourceMonad.join(ResourceMonad.map(a, UnstructuredResourceMonad::distribute)), UnstructuredMonad::join);
     // }
 
+    /**
+     * Efficiently reduce a collection of unstructured resources using an operator on their values.
+     *
+     * @see UnstructuredResourceApplicative#reduce(Collection, Object, BiFunction, String)
+     * @see ResourceMonad#reduce(Collection, Object, BiFunction)
+     */
+    public static <A> Resource<Unstructured<A>> reduce(Collection<? extends Resource<Unstructured<A>>> operands, A identity, BiFunction<A, A, A> f) {
+        return ResourceMonad.reduce(operands, UnstructuredMonad.pure(identity), UnstructuredMonad.map(f));
+    }
+
+    /**
+     * Like {@link UnstructuredResourceApplicative#reduce(Collection, Object, BiFunction)}, but names the result.
+     *
+     * @see UnstructuredResourceApplicative#reduce(Collection, Object, BiFunction)
+     * @see ResourceMonad#reduce(Collection, Object, BiFunction, String)
+     */
+    public static <A> Resource<Unstructured<A>> reduce(Collection<? extends Resource<Unstructured<A>>> operands, A identity, BiFunction<A, A, A> f, String operationName) {
+        return ResourceMonad.reduce(operands, UnstructuredMonad.pure(identity), UnstructuredMonad.map(f), operationName);
+    }
+
     // GENERATED CODE START
     // Supplemental methods generated by generate_monad_methods.py on 2023-12-06.
 

diff --git a/.../main/java/gov/nasa/jpl/aerie/contrib/streamline/modeling/discrete/DiscreteResources.java b/.../main/java/gov/nasa/jpl/aerie/contrib/streamline/modeling/discrete/DiscreteResources.java
@@ -18,7 +18,6 @@
 import java.util.*;
 import java.util.function.BiPredicate;
 import java.util.function.Supplier;
-import java.util.stream.Stream;
 
 import static gov.nasa.jpl.aerie.contrib.streamline.core.CellRefV2.autoEffects;
 import static gov.nasa.jpl.aerie.contrib.streamline.core.CellRefV2.testing;
@@ -199,15 +198,17 @@ public static Resource<Discrete<Boolean>> and(Resource<Discrete<Boolean>> left,
    */
   @SafeVarargs
   public static Resource<Discrete<Boolean>> all(Resource<Discrete<Boolean>>... operands) {
-    return all(stream(operands));
+    return all(stream(operands).toList());
   }
 
   /**
    * Reduce operands using short-circuiting logical "and"
    */
-  public static Resource<Discrete<Boolean>> all(Stream<? extends Resource<Discrete<Boolean>>> operands) {
+  public static Resource<Discrete<Boolean>> all(Collection<? extends Resource<Discrete<Boolean>>> operands) {
     // Reduce using the short-circuiting and to improve efficiency
-    return reduce(operands, constant(true), DiscreteResources::and, "All");
+    // Unlike most reductions, we explicitly want to add intermediate nodes by using Resources.reduce instead of ResourceMonad.reduce.
+    // Those intermediate nodes allow us to short-circuit, truncating our dependencies.
+    return Resources.reduce(operands, constant(true), DiscreteResources::and, "All");
   }
 
   /**
@@ -225,15 +226,17 @@ public static Resource<Discrete<Boolean>> or(Resource<Discrete<Boolean>> left, R
    */
   @SafeVarargs
   public static Resource<Discrete<Boolean>> any(Resource<Discrete<Boolean>>... operands) {
-    return any(stream(operands));
+    return any(stream(operands).toList());
   }
 
   /**
    * Reduce operands using short-circuiting logical "or"
    */
-  public static Resource<Discrete<Boolean>> any(Stream<? extends Resource<Discrete<Boolean>>> operands) {
+  public static Resource<Discrete<Boolean>> any(Collection<? extends Resource<Discrete<Boolean>>> operands) {
     // Reduce using the short-circuiting or to improve efficiency
-    return reduce(operands, constant(false), DiscreteResources::or, "Any");
+    // Unlike most reductions, we explicitly want to add intermediate nodes by using Resources.reduce instead of ResourceMonad.reduce.
+    // Those intermediate nodes allow us to short-circuit, truncating our dependencies.
+    return Resources.reduce(operands, constant(false), DiscreteResources::or, "Any");
   }
 
   /**
@@ -278,14 +281,14 @@ public static Resource<Discrete<Boolean>> assertThat(String description, Resourc
    */
   @SafeVarargs
   public static Resource<Discrete<Integer>> addInt(Resource<Discrete<Integer>>... operands) {
-    return sumInt(Arrays.stream(operands));
+    return sumInt(Arrays.stream(operands).toList());
   }
 
   /**
    * Add integer resources
    */
-  public static Resource<Discrete<Integer>> sumInt(Stream<? extends Resource<Discrete<Integer>>> operands) {
-    return reduce(operands, constant(0), map(Integer::sum), "Sum");
+  public static Resource<Discrete<Integer>> sumInt(Collection<? extends Resource<Discrete<Integer>>> operands) {
+    return reduce(operands, 0, Integer::sum, "Sum");
   }
 
   /**
@@ -302,14 +305,14 @@ public static Resource<Discrete<Integer>> subtractInt(Resource<Discrete<Integer>
    */
   @SafeVarargs
   public static Resource<Discrete<Integer>> multiplyInt(Resource<Discrete<Integer>>... operands) {
-    return productInt(Arrays.stream(operands));
+    return productInt(Arrays.stream(operands).toList());
   }
 
   /**
    * Multiply integer resources
    */
-  public static Resource<Discrete<Integer>> productInt(Stream<? extends Resource<Discrete<Integer>>> operands) {
-    return reduce(operands, constant(1), map((x, y) -> x * y), "Product");
+  public static Resource<Discrete<Integer>> productInt(Collection<? extends Resource<Discrete<Integer>>> operands) {
+    return reduce(operands, 1, (x, y) -> x * y, "Product");
   }
 
   /**
@@ -328,14 +331,14 @@ public static Resource<Discrete<Integer>> divideInt(Resource<Discrete<Integer>>
    */
   @SafeVarargs
   public static Resource<Discrete<Double>> add(Resource<Discrete<Double>>... operands) {
-    return sum(Arrays.stream(operands));
+    return sum(Arrays.stream(operands).toList());
   }
 
   /**
    * Add double resources
    */
-  public static Resource<Discrete<Double>> sum(Stream<? extends Resource<Discrete<Double>>> operands) {
-    return reduce(operands, constant(0.0), map(Double::sum), "Sum");
+  public static Resource<Discrete<Double>> sum(Collection<? extends Resource<Discrete<Double>>> operands) {
+    return reduce(operands, 0.0, Double::sum, "Sum");
   }
 
   /**
@@ -352,14 +355,14 @@ public static Resource<Discrete<Double>> subtract(Resource<Discrete<Double>> lef
    */
   @SafeVarargs
   public static Resource<Discrete<Double>> multiply(Resource<Discrete<Double>>... operands) {
-    return product(Arrays.stream(operands));
+    return product(Arrays.stream(operands).toList());
   }
 
   /**
    * Multiply double resources
    */
-  public static Resource<Discrete<Double>> product(Stream<? extends Resource<Discrete<Double>>> operands) {
-    return reduce(operands, constant(1.0), map((x, y) -> x * y), "Product");
+  public static Resource<Discrete<Double>> product(Collection<? extends Resource<Discrete<Double>>> operands) {
+    return reduce(operands, 1.0, (x, y) -> x * y, "Product");
   }
 
   /**

diff --git a/...gov/nasa/jpl/aerie/contrib/streamline/modeling/discrete/monads/DiscreteResourceMonad.java b/...gov/nasa/jpl/aerie/contrib/streamline/modeling/discrete/monads/DiscreteResourceMonad.java
@@ -6,6 +6,7 @@
 import gov.nasa.jpl.aerie.contrib.streamline.utils.*;
 import org.apache.commons.lang3.function.TriFunction;
 
+import java.util.Collection;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -30,6 +31,26 @@ public static <A> Resource<Discrete<A>> join(Resource<Discrete<Resource<Discrete
     return ResourceMonad.map(ResourceMonad.join(ResourceMonad.map(a, DiscreteResourceMonad::distribute)), DiscreteMonad::join);
   }
 
+  /**
+   * Efficiently reduce a collection of discrete resources using an operator on their values.
+   *
+   * @see DiscreteResourceMonad#reduce(Collection, Object, BiFunction, String)
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction)
+   */
+  public static <A> Resource<Discrete<A>> reduce(Collection<? extends Resource<Discrete<A>>> operands, A identity, BiFunction<A, A, A> f) {
+    return ResourceMonad.reduce(operands, DiscreteMonad.pure(identity), DiscreteMonad.map(f));
+  }
+
+  /**
+   * Like {@link DiscreteResourceMonad#reduce(Collection, Object, BiFunction)}, but names the result.
+   *
+   * @see DiscreteResourceMonad#reduce(Collection, Object, BiFunction)
+   * @see ResourceMonad#reduce(Collection, Object, BiFunction, String)
+   */
+  public static <A> Resource<Discrete<A>> reduce(Collection<? extends Resource<Discrete<A>>> operands, A identity, BiFunction<A, A, A> f, String operationName) {
+    return ResourceMonad.reduce(operands, DiscreteMonad.pure(identity), DiscreteMonad.map(f), operationName);
+  }
+
   // GENERATED CODE START
   // Supplemental methods generated by generate_monad_methods.py on 2023-12-06.