Add expansion population function

lphy-base/src/main/java/lphy/base/evolution/coalescent/populationmodel/ConstantPopulation.java

@@ -1,10 +1,6 @@
 package lphy.base.evolution.coalescent.populationmodel;
-import lphy.base.evolution.coalescent.PopulationFunction;
 
-import java.io.FileWriter;
-import java.io.IOException;
-import java.io.PrintWriter;
-import java.util.Locale;
+import lphy.base.evolution.coalescent.PopulationFunction;
 
 public class ConstantPopulation implements PopulationFunction {
 
@@ -62,24 +58,4 @@ public String toString() {
         return "Constant population size of " + N0;
     }
 
-    public static void main(String[] args) {
-        double N = 100;
-        double tStart = 0;
-        double tEnd = 100;
-        int nPoints = 100;
-
-        ConstantPopulation constantPopulation = new ConstantPopulation(N);
-
-        try (PrintWriter writer = new PrintWriter(new FileWriter("constant_data.csv"))) {
-            writer.println("time,theta");
-            for (int i = 0; i < nPoints; i++) {
-                double t = tStart + (i / (double) (nPoints - 1)) * (tEnd - tStart);
-                double theta = constantPopulation.getTheta(t);
-
-                writer.printf(Locale.US, "%.4f,%.4f%n", t, theta);
-            }
-        } catch (IOException e) {
-            e.printStackTrace();
-        }
-    }
 }

lphy-base/src/main/java/lphy/base/evolution/coalescent/populationmodel/ExpansionPopulation.java

@@ -0,0 +1,174 @@
+package lphy.base.evolution.coalescent.populationmodel;
+
+import lphy.base.evolution.coalescent.PopulationFunction;
+import org.apache.commons.math3.analysis.UnivariateFunction;
+import org.apache.commons.math3.analysis.integration.IterativeLegendreGaussIntegrator;
+import org.apache.commons.math3.analysis.solvers.BrentSolver;
+import org.apache.commons.math3.analysis.solvers.UnivariateSolver;
+import org.apache.commons.math3.exception.NoBracketingException;
+import org.apache.commons.math3.exception.TooManyEvaluationsException;
+
+/**
+ * ExpansionPopulation models population size changes over time following a constant-exponential model.
+ * It specifically handles scenarios with an ancestral population size (NA).
+ */
+public class ExpansionPopulation implements PopulationFunction {
+
+    private final double NC;   // Initial population size at t=0
+    private final double NA;   // Ancestral population size after decay
+    private final double r;    // Exponential decay rate
+    private final double x;    // Time at which exponential decay starts
+
+    /**
+     * Constructor for the ExpansionPopulation model with ancestral population size (NA).
+     *
+     * @param NC Initial population size at time t=0.
+     * @param NA Ancestral population size after decay.
+     * @param r  Exponential decay rate.
+     * @param x  Time at which exponential decay starts.
+     */
+    public ExpansionPopulation(double NC, double NA, double r, double x) {
+        if (NC <= 0) {
+            throw new IllegalArgumentException("Initial population size NC must be positive.");
+        }
+        if (NA <= 0) {
+            throw new IllegalArgumentException("Ancestral population size NA must be positive.");
+        }
+        if (NA >= NC) {
+            throw new IllegalArgumentException("Ancestral population size NA must be less than initial population size NC.");
+        }
+        if (r <= 0) {
+            throw new IllegalArgumentException("Decay rate r must be positive.");
+        }
+        if (x < 0) {
+            throw new IllegalArgumentException("Time x must be non-negative.");
+        }
+        this.NC = NC;
+        this.NA = NA;
+        this.r = r;
+        this.x = x;
+    }
+
+    /**
+     * Calculate the population size N(t) at a given time t.
+     *
+     * @param t Time.
+     * @return Population size at time t.
+     */
+    @Override
+    public double getTheta(double t) {
+        if (t < 0) {
+            throw new IllegalArgumentException("Time t cannot be negative.");
+        }
+
+        if (t <= x) {
+            return NC;
+        } else {
+            return (NC - NA) * Math.exp(-r * (t - x)) + NA;
+        }
+    }
+
+    /**
+     * Calculate the cumulative intensity over the time period from 0 to t.
+     *
+     * @param t Time.
+     * @return Intensity value.
+     */
+    @Override
+    public double getIntensity(double t) {
+        if (t < 0) return 0.0;
+
+        if (t <= x) {
+            return t / NC;
+        } else {
+            // Integral from 0 to x: ∫(1 / NC) dt = t / NC
+            double firstIntegral = x / NC;
+
+            // Integral from x to t: ∫(1 / [(NC - NA)e^{-r(t' - x)} + NA] ) dt
+            // No closed-form solution; use numerical integration
+
+            UnivariateFunction integrand = timePoint -> 1.0 / ((NC - NA) * Math.exp(-r * (timePoint - x)) + NA);
+            IterativeLegendreGaussIntegrator integrator = new IterativeLegendreGaussIntegrator(
+                    5, 1.0e-12, 1.0e-8, 2, 10000);
+            double secondIntegral = integrator.integrate(Integer.MAX_VALUE, integrand, x, t);
+
+            return firstIntegral + secondIntegral;
+        }
+    }
+
+    /**
+     * Calculate the inverse intensity, i.e., find time t such that Intensity(t) = x.
+     *
+     * @param x Intensity value.
+     * @return Time t corresponding to the given intensity.
+     */
+    @Override
+    public double getInverseIntensity(double x) {
+        if (x < 0) {
+            throw new IllegalArgumentException("Intensity x must be non-negative.");
+        }
+
+        UnivariateFunction equation = tPoint -> getIntensity(tPoint) - x;
+        UnivariateSolver solver = new BrentSolver(1e-9, 1e-9);
+        double tMin = 0.0;
+        double tMax = 1e6; // Arbitrary large number to ensure convergence
+
+        try {
+            return solver.solve(1000, equation, tMin, tMax);
+        } catch (NoBracketingException | TooManyEvaluationsException e) {
+            throw new IllegalArgumentException("Cannot find corresponding time t for the given intensity x.", e);
+        }
+    }
+
+    @Override
+    public boolean isAnalytical() {
+        return false;
+    }
+
+    /**
+     * Provides a string representation of the ExpansionPopulation model.
+     *
+     * @return String describing the model parameters.
+     */
+    @Override
+    public String toString() {
+        return "ExpansionPopulation [NC=" + NC + ", NA=" + NA +
+                ", r=" + r + ", x=" + x + "]";
+    }
+
+    /**
+     * Get the initial population size NC.
+     *
+     * @return Initial population size NC.
+     */
+    public double getNC() {
+        return NC;
+    }
+
+    /**
+     * Get the ancestral population size NA.
+     *
+     * @return Ancestral population size NA.
+     */
+    public double getNA() {
+        return NA;
+    }
+
+    /**
+     * Get the exponential decay rate r.
+     *
+     * @return Exponential decay rate r.
+     */
+    public double getR() {
+        return r;
+    }
+
+    /**
+     * Get the time x at which exponential decay starts.
+     *
+     * @return Time x.
+     */
+    public double getX() {
+        return x;
+    }
+}

lphy-base/src/main/java/lphy/base/evolution/coalescent/populationmodel/ExpansionPopulationFunction.java

@@ -0,0 +1,73 @@
+package lphy.base.evolution.coalescent.populationmodel;
+
+import lphy.base.evolution.coalescent.PopulationFunction;
+import lphy.core.model.DeterministicFunction;
+import lphy.core.model.Value;
+import lphy.core.model.annotation.GeneratorCategory;
+import lphy.core.model.annotation.GeneratorInfo;
+import lphy.core.model.annotation.ParameterInfo;
+
+public class ExpansionPopulationFunction extends DeterministicFunction<PopulationFunction> {
+
+    public static final String NA_PARAM_NAME = "NA";
+    public static final String R_PARAM_NAME = "r";
+    public static final String NC_PARAM_NAME = "NC";
+    public static final String X_PARAM_NAME = "x"; // New independent parameter for time x
+
+    public ExpansionPopulationFunction(
+            @ParameterInfo(name = NA_PARAM_NAME, description = "Ancestral population size after decay.")
+            Value<Double> NA,
+            @ParameterInfo(name = R_PARAM_NAME, description = "The exponential growth rate.")
+            Value<Double> r,
+            @ParameterInfo(name = NC_PARAM_NAME, description = "Current population size after time x.")
+            Value<Double> NC,
+            @ParameterInfo(name = X_PARAM_NAME, description = "Time at which the population reaches NC.")
+            Value<Double> x) {
+
+        setParam(NA_PARAM_NAME, NA);
+        setParam(R_PARAM_NAME, r);
+        setParam(NC_PARAM_NAME, NC);
+        setParam(X_PARAM_NAME, x);
+    }
+
+    @GeneratorInfo(
+            name = "ExpansionPopFunc",
+            narrativeName = "Expansion Population Function",
+            category = GeneratorCategory.COAL_TREE,
+            examples = {"expansionCoal.lphy"},
+            description = "Models population using a constant-exponential model with ancestral population size NA."
+    )
+    @Override
+    public Value<PopulationFunction> apply() {
+        Value<Double> NAValue = (Value<Double>) getParams().get(NA_PARAM_NAME);
+        Value<Double> rValue = (Value<Double>) getParams().get(R_PARAM_NAME);
+        Value<Double> NCValue = (Value<Double>) getParams().get(NC_PARAM_NAME);
+        Value<Double> xValue = (Value<Double>) getParams().get(X_PARAM_NAME); // Use x parameter
+
+        double NA = NAValue.value();
+        double r = rValue.value();
+        double NC = NCValue.value();
+        double x = xValue.value();
+
+        PopulationFunction expansionPopulation = new ExpansionPopulation(NC, NA, r, x);
+
+
+        return new Value<>(expansionPopulation, this);
+    }
+
+    public Value<Double> getNA() {
+        return (Value<Double>) getParams().get(NA_PARAM_NAME);
+    }
+
+    public Value<Double> getR() {
+        return (Value<Double>) getParams().get(R_PARAM_NAME);
+    }
+
+    public Value<Double> getNC() {
+        return (Value<Double>) getParams().get(NC_PARAM_NAME);
+    }
+
+    public Value<Double> getX() {
+        return (Value<Double>) getParams().get(X_PARAM_NAME); // Getter for x parameter
+    }
+}

lphy-base/src/main/java/lphy/base/evolution/coalescent/populationmodel/ExponentialPopulation.java

@@ -6,8 +6,6 @@
 import org.apache.commons.math3.analysis.solvers.BrentSolver;
 import org.apache.commons.math3.analysis.solvers.UnivariateSolver;
 
-import java.util.Locale;
-
 public class ExponentialPopulation implements PopulationFunction {
 
     private double growthRate;
@@ -210,40 +208,5 @@ public String toString() {
                     ", Initial Size = " + N0;
         }
     }
-
-    /**
-     * Main method for testing the ExponentialPopulation class functionality.
-     *
-     * @param args Command-line arguments.
-     */
-    public static void main(String[] args) {
-        double growthRate = 0.1;
-        double N0 = 100.0;
-        double ancestralPopulation = 50.0;
-
-        // Define specific test time points
-        double[] testTimes = {0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0};
-
-        // Create an instance without NA
-        ExponentialPopulation expPop = new ExponentialPopulation(growthRate, N0);
-
-        // Create an instance with NA
-        ExponentialPopulation expPopWithNA = new ExponentialPopulation(growthRate, N0, ancestralPopulation);
-
-        // Print header for test points
-        System.out.println("time,theta_noNA,theta_withNA,intensity_noNA,intensity_withNA");
-
-        // Iterate over each test time point
-        for (double t : testTimes) {
-            double theta_noNA = expPop.getTheta(t);
-            double theta_withNA = expPopWithNA.getTheta(t);
-            double intensity_noNA = expPop.getIntensity(t);
-            double intensity_withNA = expPopWithNA.getIntensity(t);
-
-            // Print the results in CSV format
-            System.out.printf(Locale.US, "%.2f,%.4f,%.4f,%.6f,%.6f%n",
-                    t, theta_noNA, theta_withNA, intensity_noNA, intensity_withNA);
-        }
-    }
 }
 

lphy-base/src/main/java/lphy/base/spi/LPhyBaseImpl.java

@@ -113,7 +113,7 @@ public List<Class<? extends BasicFunction>> declareFunctions() {
                 GompertzPopulationFunction_f0.class,
                 GompertzPopulationFunction_t50.class, ExponentialPopulationFunction.class,
                 LogisticPopulationFunction.class, ConstantPopulationFunction.class,
-                Cons_Exp_ConsPopulationFunction.class
+                Cons_Exp_ConsPopulationFunction.class, ExpansionPopulationFunction.class
 
         );