verify units for probability parameters in math generation

virtualcell · Sep 24, 2024 · 15200c2 · 15200c2
1 parent 6f97352
commit 15200c2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/vcell-core/src/main/java/cbit/vcell/mapping/AbstractMathMapping.java b/vcell-core/src/main/java/cbit/vcell/mapping/AbstractMathMapping.java
@@ -62,6 +62,7 @@
 import cbit.vcell.solver.Simulation;
 import cbit.vcell.units.VCUnitDefinition;
 import cbit.vcell.units.VCUnitException;
+import ucar.units_vcell.UnitException;
 
 public abstract class AbstractMathMapping implements ScopedSymbolTable, UnitFactorProvider, IssueSource, MathMapping {
 
@@ -1749,13 +1750,24 @@ protected final SpeciesCountParameter addSpeciesCountParameter(String name, Expr
      */
     final ProbabilityParameter addProbabilityParameter(String name, Expression expression, int role,
                                                        VCUnitDefinition unitDefinition, ModelProcess argModelProcess) throws java.beans.PropertyVetoException,
-            ExpressionBindingException{
+            VCUnitException {
 
         GeometryClass geometryClass = null;
-        if(argModelProcess.getStructure() != null){
+        if (argModelProcess.getStructure() != null){
             geometryClass = simContext.getGeometryContext().getStructureMapping(argModelProcess.getStructure()).getGeometryClass();
         }
         ProbabilityParameter newParameter = new ProbabilityParameter(name, expression, role, unitDefinition, argModelProcess, geometryClass);
+        try {
+            VCUnitEvaluator unitEvaluator = new VCUnitEvaluator(simContext.getModel().getUnitSystem());
+            VCUnitDefinition expUnit = unitEvaluator.getUnitDefinition(expression);
+            if (unitDefinition != null && !unitDefinition.equals(expUnit)){
+                logger.error("expected unit=" + unitDefinition.getSymbol() + ", found=" + expUnit.getSymbol());
+                throw new VCUnitException("expected unit=" + unitDefinition.getSymbol() + ", found=" + expUnit.getSymbol());
+            }
+        } catch (ExpressionException e){
+            logger.error("error evaluating units for expression '" + expression + ": " + e.getMessage(), e);
+            throw new VCUnitException("error evaluating units for expression '" + expression + ": " + e.getMessage());
+        }
         MathMappingParameter previousParameter = getMathMappingParameter(name);
         if(previousParameter != null){
             logger.info("MathMappingParameter addProbabilityParameter found duplicate parameter for name " + name);

diff --git a/vcell-core/src/main/java/cbit/vcell/mapping/StochMathMapping.java b/vcell-core/src/main/java/cbit/vcell/mapping/StochMathMapping.java
@@ -144,7 +144,7 @@ private Expression getProbabilityRate(ReactionStep reactionStep, MassActionStoch
 			Expression massActionRateCoefficient = isForwardDirection ? massActionStochasticFunction.forwardRate() : massActionStochasticFunction.reverseRate();
 			probExp = Expression.mult(massActionRateCoefficient, productOfSpeciesCounts, factorsExp);
 		}
-		return probExp;
+		return probExp.flattenSafe();
 	}
 
 	private Expression getProbabilityRate(ReactionStep reactionStep, GeneralKineticsStochasticFunction generalKineticsStochasticFunction, boolean isForwardDirection)
@@ -188,7 +188,7 @@ private Expression getProbabilityRate(ReactionStep reactionStep, GeneralKinetics
 		simplifiedFactorExp.bindExpression(this);
 		//get probability rate with converting factor
 		Expression probExp = Expression.mult(netRateExpr, simplifiedFactorExp);
-		probExp = probExp.flatten();
+		probExp = probExp.flattenSafe();
 		return probExp;
 	}