Merge pull request openrocket#2658 from SiboVG/issue-2657

[openrocket#2657] Do not ignore active child stages of inactive parent stage in calculations

core/src/main/java/info/openrocket/core/aerodynamics/BarrowmanCalculator.java

@@ -188,6 +188,18 @@ private AerodynamicForces calculateForceAnalysis(FlightConfiguration configurati
 		for (RocketComponent child : comp.getChildren()) {
 			// Ignore inactive stages
 			if (child instanceof AxialStage && !configuration.isStageActive(child.getStageNumber())) {
+				// Check if there are child stages that are active and need to be calculated
+				for (AxialStage childStage : child.getTopLevelChildStages()) {
+					if (configuration.isStageActive(childStage.getStageNumber())) {
+						// forces particular to each component
+						AerodynamicForces childForces = calculateForceAnalysis(configuration, conds, childStage, instances,
+								eachForces, assemblyForces, warnings);
+
+						if (childForces != null) {
+							aggregateForces.merge(childForces);
+						}
+					}
+				}
 				continue;
 			}
 
@@ -308,13 +320,8 @@ private AerodynamicForces calculateNonAxialForces(FlightConfiguration configurat
 	public void checkGeometry(FlightConfiguration configuration, final RocketComponent treeRoot, WarningSet warnings) {
 		Queue<RocketComponent> queue = new LinkedList<>();
 
-		for (RocketComponent child : treeRoot.getChildren()) {
-			// Ignore inactive stages
-			if (child instanceof AxialStage && !configuration.isStageActive(child.getStageNumber())) {
-				continue;
-			}
-			queue.add(child);
-		}
+		// Add the (active) child stages
+		addDirectChildStagesToQueue(configuration, queue, treeRoot);
 
 		SymmetricComponent prevComp = null;
 		if ((treeRoot instanceof ComponentAssembly) &&
@@ -328,13 +335,8 @@ public void checkGeometry(FlightConfiguration configuration, final RocketCompone
 			if ((comp instanceof SymmetricComponent) ||
 					((comp instanceof AxialStage) &&
 							!(comp instanceof ParallelStage))) {
-				for (RocketComponent child : comp.getChildren()) {
-					// Ignore inactive stages
-					if (child instanceof AxialStage && !configuration.isStageActive(child.getStageNumber())) {
-						continue;
-					}
-					queue.add(child);
-				}
+				// Add the (active) child stages
+				addDirectChildStagesToQueue(configuration, queue, comp);
 
 				if (comp instanceof SymmetricComponent) {
 					SymmetricComponent sym = (SymmetricComponent) comp;
@@ -440,6 +442,29 @@ public void checkGeometry(FlightConfiguration configuration, final RocketCompone
 		}
 	}
 
+	/**
+	 * Add child stages to the queue. Only active stages are added. If a child stages is inactive, but it does have
+	 * active child stages, these are added to the queue.
+	 * @param configuration Rocket configuration
+	 * @param queue Queue to add stages to
+	 * @param comp Component to add stages from
+	 */
+	private void addDirectChildStagesToQueue(FlightConfiguration configuration, Queue<RocketComponent> queue, RocketComponent comp) {
+		for (RocketComponent child : comp.getChildren()) {
+			// Ignore inactive stages
+			if (child instanceof AxialStage && !configuration.isStageActive(child.getStageNumber())) {
+				// Check if there are child stages that are active and need to be calculated
+				for (AxialStage childStage : child.getTopLevelChildStages()) {
+					if (configuration.isStageActive(childStage.getStageNumber())) {
+						queue.add(childStage);
+					}
+				}
+				continue;
+			}
+			queue.add(child);
+		}
+	}
+
 	//////////////// DRAG CALCULATIONS ////////////////
 	/**
 	 * Calculation of drag coefficient due to air friction
@@ -540,6 +565,9 @@ private double calculateFrictionCD(FlightConfiguration configuration, FlightCond
 			}
 
 			if (forceMap != null) {
+				if (forceMap.get(c) == null) {
+					System.out.println("No forces for " + c);
+				}
 				forceMap.get(c).setFrictionCD(componentFrictionCD);
 			}
 		}

core/src/main/java/info/openrocket/core/rocketcomponent/RocketComponent.java

@@ -2365,6 +2365,54 @@ public final List<ComponentAssembly> getDirectChildAssemblies() {
 		return result;
 	}
 
+	/**
+	 * Return all the stages that are a child of this component.
+	 * @return all the stages that are a child of this component.
+	 */
+	public final List<AxialStage> getAllChildStages() {
+		checkState();
+
+		Iterator<RocketComponent> children = iterator(false);
+
+		List<AxialStage> result = new ArrayList<>();
+
+		while (children.hasNext()) {
+			RocketComponent child = children.next();
+			if (child instanceof AxialStage) {
+				result.add((AxialStage) child);
+			}
+		}
+		return result;
+	}
+
+	/**
+	 * Return all the stages that are a child of this component without counting child stages of the found stages.
+	 * @return all the stages that are a child of this component without counting child stages of the found stages.
+	 */
+	public final List<AxialStage> getTopLevelChildStages() {
+		checkState();
+
+		List<AxialStage> result = new ArrayList<>();
+		addTopLevelStagesToList(result, this);
+
+		return result;
+	}
+
+	/**
+	 * Add all the top-level stages of the given component to the list.
+	 * @param list list to add the top-level stages to
+	 * @param parent parent component to search for top-level stages
+	 */
+	private void addTopLevelStagesToList(List<AxialStage> list, RocketComponent parent) {
+		for (RocketComponent child : parent.getChildren()) {
+			if (child instanceof AxialStage) {
+				list.add((AxialStage) child);
+			} else {
+				addTopLevelStagesToList(list, child);
+			}
+		}
+	}
+
 	/**
 	 * Return all the component assemblies that are a parent or super-parent of this component
 	 * @return list of ComponentAssembly components that are a parent or super-parent of this component

core/src/test/java/info/openrocket/core/aerodynamics/BarrowmanCalculatorTest.java

@@ -2,9 +2,14 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import info.openrocket.core.document.Simulation;
 import info.openrocket.core.logging.WarningSet;
+import info.openrocket.core.rocketcomponent.ComponentAssembly;
+import info.openrocket.core.rocketcomponent.FlightConfigurationId;
+import info.openrocket.core.rocketcomponent.RocketComponent;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
@@ -30,6 +35,8 @@
 import info.openrocket.core.util.MathUtil;
 import info.openrocket.core.util.TestRockets;
 
+import java.util.Map;
+
 public class BarrowmanCalculatorTest {
 	protected final double EPSILON = 0.00001;
 
@@ -664,4 +671,99 @@ public void testRailButtonDrag() {
 		final double zeroCD = zeroForces.getCD();
 		assertEquals(epsCD, zeroCD, EPSILON, "drag at mach 0 should equal drag at mach MathUtil.EPSILON");
 	}
+
+	/**
+	 * Tests that the force map contains all necessary components when stages are enabled/disabled.
+	 */
+	@Test
+	public void testForceMapContent() {
+		Rocket rocket = TestRockets.makeFalcon9Heavy();
+		FlightConfigurationId fcid = new FlightConfigurationId(TestRockets.FALCON_9H_FCID_1);
+
+		AerodynamicCalculator calculator = new BarrowmanCalculator();
+
+		Simulation simulation = new Simulation(rocket);
+		simulation.setFlightConfigurationId(fcid);
+		FlightConfiguration config = simulation.getActiveConfiguration();
+		FlightConditions conditions = new FlightConditions(config);
+		WarningSet warnings = new WarningSet();
+
+		// Get force analysis with all stages active
+		Map<RocketComponent, AerodynamicForces> forceMap = calculator.getForceAnalysis(config, conditions, warnings);
+
+		// First verify all stages active case
+		assertTrue(config.isStageActive(TestRockets.FALCON_9H_PAYLOAD_STAGE_NUMBER));
+		assertTrue(config.isStageActive(TestRockets.FALCON_9H_CORE_STAGE_NUMBER));
+		assertTrue(config.isStageActive(TestRockets.FALCON_9H_BOOSTER_STAGE_NUMBER));
+
+		// Check that force map contains all aerodynamic components
+		for (RocketComponent comp : rocket) {
+			if (comp.isAerodynamic() || comp instanceof ComponentAssembly) {
+				assertTrue(forceMap.containsKey(comp),
+						"Force map missing component: " + comp.getName());
+				assertNotNull(forceMap.get(comp),
+						"Force entry is null for: " + comp.getName());
+			}
+		}
+		assertEquals(13, forceMap.size(), "Force map should contain 13 components");
+
+		// Now disable core stage and check map is updated correctly
+		config._setStageActive(TestRockets.FALCON_9H_CORE_STAGE_NUMBER, false);
+		config._setStageActive(TestRockets.FALCON_9H_BOOSTER_STAGE_NUMBER, false);
+
+		warnings.clear();
+		Map<RocketComponent, AerodynamicForces> disabledForceMap = calculator.getForceAnalysis(config, conditions, warnings);
+
+		// Map should only contain components from active stages
+		for (RocketComponent comp : rocket) {
+			if (!config.isComponentActive(comp)) {
+				assertFalse(disabledForceMap.containsKey(comp),
+						"Force map contains inactive component: " + comp.getName());
+			}
+			else if (comp.isAerodynamic() || comp instanceof ComponentAssembly) {
+				assertTrue(disabledForceMap.containsKey(comp),
+						"Force map missing active component: " + comp.getName());
+				assertNotNull(disabledForceMap.get(comp),
+						"Force entry is null for active component: " + comp.getName());
+			}
+		}
+		assertEquals(7, disabledForceMap.size(), "Force map should contain 7 components");
+
+		// Re-enable the stages and verify map is complete again
+		config.setAllStages();
+
+		warnings.clear();
+		Map<RocketComponent, AerodynamicForces> reenabledForceMap = calculator.getForceAnalysis(config, conditions, warnings);
+
+		for (RocketComponent comp : rocket) {
+			if (comp.isAerodynamic() || comp instanceof ComponentAssembly) {
+				assertTrue(reenabledForceMap.containsKey(comp),
+						"Force map missing component after re-enable: " + comp.getName());
+				assertNotNull(reenabledForceMap.get(comp),
+						"Force entry is null after re-enable for: " + comp.getName());
+			}
+		}
+
+		// Now disable the core stage, but enable its child stage, the booster
+		config._setStageActive(TestRockets.FALCON_9H_CORE_STAGE_NUMBER, false);
+		config._setStageActive(TestRockets.FALCON_9H_BOOSTER_STAGE_NUMBER, true);
+
+		warnings.clear();
+		Map<RocketComponent, AerodynamicForces> boosterOnlyForceMap = calculator.getForceAnalysis(config, conditions, warnings);
+
+		// Map should only contain components from active stages
+		for (RocketComponent comp : rocket) {
+			if (!config.isComponentActive(comp)) {
+				assertFalse(boosterOnlyForceMap.containsKey(comp),
+						"Force map contains inactive component: " + comp.getName());
+			}
+			else if (comp.isAerodynamic() || comp instanceof ComponentAssembly) {
+				assertTrue(boosterOnlyForceMap.containsKey(comp),
+						"Force map missing active component: " + comp.getName());
+				assertNotNull(boosterOnlyForceMap.get(comp),
+						"Force entry is null for active component: " + comp.getName());
+			}
+		}
+		assertEquals(11, boosterOnlyForceMap.size(), "Force map should contain 10 components");
+	}
 }