Thrust calculated with minThrust
[VesselSimulator.git] / KerbalEngineer / VesselSimulator / EngineSim.cs
blob:a/KerbalEngineer/VesselSimulator/EngineSim.cs -> blob:b/KerbalEngineer/VesselSimulator/EngineSim.cs
--- a/KerbalEngineer/VesselSimulator/EngineSim.cs
+++ b/KerbalEngineer/VesselSimulator/EngineSim.cs
@@ -1,364 +1,483 @@
-// Kerbal Engineer Redux

-// Author:  CYBUTEK

-// License: Attribution-NonCommercial-ShareAlike 3.0 Unported

-

-#region

-

-using System;

-using System.Collections.Generic;

-using System.Linq;

-using System.Text;

-

-using UnityEngine;

-

-#endregion

-

-namespace KerbalEngineer.VesselSimulator

-{

-    public class EngineSim

-    {

-        private readonly ResourceContainer resourceConsumptions = new ResourceContainer();

-

-        public double actualThrust = 0;

-        public bool isActive = false;

-        public double isp = 0;

-        public PartSim partSim;

-

-        public double thrust = 0;

-

-        // Add thrust vector to account for directional losses

-        public Vector3 thrustVec;

-

-        public EngineSim(PartSim theEngine,

-                         double atmosphere,

-                         double velocity,

-                         float maxThrust,

-                         float minThrust,

-                         float thrustPercentage,

-                         float requestedThrust,

-                         Vector3 vecThrust,

-                         float realIsp,

-                         FloatCurve atmosphereCurve,

-                         FloatCurve velocityCurve,

-                         bool throttleLocked,

-                         List<Propellant> propellants,

-                         bool active,

-                         bool correctThrust)

-        {

-            StringBuilder buffer = null;

-            //MonoBehaviour.print("Create EngineSim for " + theEngine.name);

-            //MonoBehaviour.print("maxThrust = " + maxThrust);

-            //MonoBehaviour.print("thrustPercentage = " + thrustPercentage);

-            //MonoBehaviour.print("requestedThrust = " + requestedThrust);

-            //MonoBehaviour.print("velocity = " + velocity);

-

-            this.partSim = theEngine;

-

-            this.isActive = active;

-            this.thrust = (maxThrust - minThrust) * (thrustPercentage / 100f) + minThrust;

-            //MonoBehaviour.print("thrust = " + thrust);

-

-            this.thrustVec = vecThrust;

-

-            double flowRate = 0d;

-            if (this.partSim.hasVessel)

-            {

-                //MonoBehaviour.print("hasVessel is true");

-                this.actualThrust = requestedThrust;

-                if (velocityCurve != null)

-                {

-                    this.actualThrust *= velocityCurve.Evaluate((float)velocity);

-                    //MonoBehaviour.print("actualThrust at velocity = " + actualThrust);

-                }

-

-                this.isp = atmosphereCurve.Evaluate((float)this.partSim.part.staticPressureAtm);

-                if (this.isp == 0d)

-                {

-                    MonoBehaviour.print("Isp at " + this.partSim.part.staticPressureAtm + " is zero. Flow rate will be NaN");

-                }

-

-                if (correctThrust && realIsp == 0)

-                {

-                    float ispsl = atmosphereCurve.Evaluate(0);

-                    if (ispsl != 0)

-                    {

-                        this.thrust = this.thrust * this.isp / ispsl;

-                    }

-                    else

-                    {

-                        MonoBehaviour.print("Isp at sea level is zero. Unable to correct thrust.");

-                    }

-                    //MonoBehaviour.print("corrected thrust = " + thrust);

-                }

-

-                if (velocityCurve != null)

-                {

-                    this.thrust *= velocityCurve.Evaluate((float)velocity);

-                    //MonoBehaviour.print("thrust at velocity = " + thrust);

-                }

-

-                if (throttleLocked)

-                {

-                    //MonoBehaviour.print("throttleLocked is true");

-                    flowRate = this.thrust / (this.isp * 9.81d);

-                }

-                else

-                {

-                    if (this.partSim.isLanded)

-                    {

-                        //MonoBehaviour.print("partSim.isLanded is true, mainThrottle = " + FlightInputHandler.state.mainThrottle);

-                        flowRate = Math.Max(0.000001d, this.thrust * FlightInputHandler.state.mainThrottle) / (this.isp * 9.81d);

-                    }

-                    else

-                    {

-                        if (requestedThrust > 0)

-                        {

-                            if (velocityCurve != null)

-                            {

-                                requestedThrust *= velocityCurve.Evaluate((float)velocity);

-                                //MonoBehaviour.print("requestedThrust at velocity = " + requestedThrust);

-                            }

-

-                            //MonoBehaviour.print("requestedThrust > 0");

-                            flowRate = requestedThrust / (this.isp * 9.81d);

-                        }

-                        else

-                        {

-                            //MonoBehaviour.print("requestedThrust <= 0");

-                            flowRate = this.thrust / (this.isp * 9.81d);

-                        }

-                    }

-                }

-            }

-            else

-            {

-                //MonoBehaviour.print("hasVessel is false");

-                this.isp = atmosphereCurve.Evaluate((float)atmosphere);

-                if (this.isp == 0d)

-                {

-                    MonoBehaviour.print("Isp at " + atmosphere + " is zero. Flow rate will be NaN");

-                }

-                if (correctThrust)

-                {

-                    float ispsl = atmosphereCurve.Evaluate(0);

-                    if (ispsl != 0)

-                    {

-                        this.thrust = this.thrust * this.isp / ispsl;

-                    }

-                    else

-                    {

-                        MonoBehaviour.print("Isp at sea level is zero. Unable to correct thrust.");

-                    }

-                    //MonoBehaviour.print("corrected thrust = " + thrust);

-                }

-

-                if (velocityCurve != null)

-                {

-                    this.thrust *= velocityCurve.Evaluate((float)velocity);

-                    //MonoBehaviour.print("thrust at velocity = " + thrust);

-                }

-

-                flowRate = this.thrust / (this.isp * 9.81d);

-            }

-

-            if (SimManager.logOutput)

-            {

-                buffer = new StringBuilder(1024);

-                buffer.AppendFormat("flowRate = {0:g6}\n", flowRate);

-            }

-

-            float flowMass = 0f;

-            foreach (Propellant propellant in propellants)

-            {

-                flowMass += propellant.ratio * ResourceContainer.GetResourceDensity(propellant.id);

-            }

-

-            if (SimManager.logOutput)

-            {

-                buffer.AppendFormat("flowMass = {0:g6}\n", flowMass);

-            }

-

-            foreach (Propellant propellant in propellants)

-            {

-                if (propellant.name == "ElectricCharge" || propellant.name == "IntakeAir")

-                {

-                    continue;

-                }

-

-                double consumptionRate = propellant.ratio * flowRate / flowMass;

-                if (SimManager.logOutput)

-                {

-                    buffer.AppendFormat("Add consumption({0}, {1}:{2:d}) = {3:g6}\n", ResourceContainer.GetResourceName(propellant.id), theEngine.name, theEngine.partId, consumptionRate);

-                }

-                this.resourceConsumptions.Add(propellant.id, consumptionRate);

-            }

-

-            if (SimManager.logOutput)

-            {

-                MonoBehaviour.print(buffer);

-            }

-        }

-

-        public ResourceContainer ResourceConsumptions

-        {

-            get { return this.resourceConsumptions; }

-        }

-

-        public bool SetResourceDrains(List<PartSim> allParts, List<PartSim> allFuelLines, HashSet<PartSim> drainingParts)

-        {

-            LogMsg log = null;

-

-            // A dictionary to hold a set of parts for each resource

-            Dictionary<int, HashSet<PartSim>> sourcePartSets = new Dictionary<int, HashSet<PartSim>>();

-

-            foreach (int type in this.resourceConsumptions.Types)

-            {

-                HashSet<PartSim> sourcePartSet = null;

-                switch (ResourceContainer.GetResourceFlowMode(type))

-                {

-                    case ResourceFlowMode.NO_FLOW:

-                        if (this.partSim.resources[type] > SimManager.RESOURCE_MIN)

-                        {

-                            sourcePartSet = new HashSet<PartSim>();

-                            //MonoBehaviour.print("SetResourceDrains(" + name + ":" + partId + ") setting sources to just this");

-                            sourcePartSet.Add(this.partSim);

-                        }

-                        break;

-

-                    case ResourceFlowMode.ALL_VESSEL:

-                        foreach (PartSim aPartSim in allParts)

-                        {

-                            if (aPartSim.resources[type] > SimManager.RESOURCE_MIN)

-                            {

-                                if (sourcePartSet == null)

-                                {

-                                    sourcePartSet = new HashSet<PartSim>();

-                                }

-

-                                sourcePartSet.Add(aPartSim);

-                            }

-                        }

-                        break;

-

-                    case ResourceFlowMode.STAGE_PRIORITY_FLOW:

-                    {

-                        Dictionary<int, HashSet<PartSim>> stagePartSets = new Dictionary<int, HashSet<PartSim>>();

-                        int maxStage = -1;

-                        foreach (PartSim aPartSim in allParts)

-                        {

-                            if (aPartSim.resources[type] > SimManager.RESOURCE_MIN)

-                            {

-                                //int stage = aPartSim.decoupledInStage;            // Use the number of the stage the tank is decoupled in

-                                int stage = aPartSim.DecouplerCount(); // Use the count of decouplers between tank and root

-                                if (stage > maxStage)

-                                {

-                                    maxStage = stage;

-                                }

-                                if (stagePartSets.ContainsKey(stage))

-                                {

-                                    sourcePartSet = stagePartSets[stage];

-                                }

-                                else

-                                {

-                                    sourcePartSet = new HashSet<PartSim>();

-                                    stagePartSets.Add(stage, sourcePartSet);

-                                }

-

-                                sourcePartSet.Add(aPartSim);

-                            }

-                        }

-

-                        while (maxStage >= 0)

-                        {

-                            if (stagePartSets.ContainsKey(maxStage))

-                            {

-                                if (stagePartSets[maxStage].Count() > 0)

-                                {

-                                    sourcePartSet = stagePartSets[maxStage];

-                                    break;

-                                }

-                            }

-                            maxStage--;

-                        }

-                    }

-                        break;

-

-                    case ResourceFlowMode.STACK_PRIORITY_SEARCH:

-                        HashSet<PartSim> visited = new HashSet<PartSim>();

-

-                        if (SimManager.logOutput)

-                        {

-                            log = new LogMsg();

-                            log.buf.AppendLine("Find " + ResourceContainer.GetResourceName(type) + " sources for " + this.partSim.name + ":" + this.partSim.partId);

-                        }

-                        sourcePartSet = this.partSim.GetSourceSet(type, allParts, visited, log, "");

-                        if (SimManager.logOutput)

-                        {

-                            MonoBehaviour.print(log.buf);

-                        }

-                        break;

-

-                    default:

-                        MonoBehaviour.print("SetResourceDrains(" + this.partSim.name + ":" + this.partSim.partId + ") Unexpected flow type for " + ResourceContainer.GetResourceName(type) + ")");

-                        break;

-                }

-

-                if (sourcePartSet != null && sourcePartSet.Count > 0)

-                {

-                    sourcePartSets[type] = sourcePartSet;

-                    if (SimManager.logOutput)

-                    {

-                        log = new LogMsg();

-                        log.buf.AppendLine("Source parts for " + ResourceContainer.GetResourceName(type) + ":");

-                        foreach (PartSim partSim in sourcePartSet)

-                        {

-                            log.buf.AppendLine(partSim.name + ":" + partSim.partId);

-                        }

-                        MonoBehaviour.print(log.buf);

-                    }

-                }

-            }

-

-            // If we don't have sources for all the needed resources then return false without setting up any drains

-            foreach (int type in this.resourceConsumptions.Types)

-            {

-                if (!sourcePartSets.ContainsKey(type))

-                {

-                    if (SimManager.logOutput)

-                    {

-                        MonoBehaviour.print("No source of " + ResourceContainer.GetResourceName(type));

-                    }

-

-                    this.isActive = false;

-                    return false;

-                }

-            }

-

-            // Now we set the drains on the members of the sets and update the draining parts set

-            foreach (int type in this.resourceConsumptions.Types)

-            {

-                HashSet<PartSim> sourcePartSet = sourcePartSets[type];

-                // Loop through the members of the set 

-                double amount = this.resourceConsumptions[type] / sourcePartSet.Count;

-                foreach (PartSim partSim in sourcePartSet)

-                {

-                    if (SimManager.logOutput)

-                    {

-                        MonoBehaviour.print("Adding drain of " + amount + " " + ResourceContainer.GetResourceName(type) + " to " + partSim.name + ":" + partSim.partId);

-                    }

-

-                    partSim.resourceDrains.Add(type, amount);

-                    drainingParts.Add(partSim);

-                }

-            }

-

-            return true;

-        }

-

-        public void DumpEngineToBuffer(StringBuilder buffer, String prefix)

-        {

-            buffer.Append(prefix);

-            buffer.AppendFormat("[thrust = {0:g6}, actual = {1:g6}, isp = {2:g6}\n", this.thrust, this.actualThrust, this.isp);

-        }

-    }

+// 
+//     Kerbal Engineer Redux
+// 
+//     Copyright (C) 2014 CYBUTEK
+// 
+//     This program is free software: you can redistribute it and/or modify
+//     it under the terms of the GNU General Public License as published by
+//     the Free Software Foundation, either version 3 of the License, or
+//     (at your option) any later version.
+// 
+//     This program is distributed in the hope that it will be useful,
+//     but WITHOUT ANY WARRANTY; without even the implied warranty of
+//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//     GNU General Public License for more details.
+// 
+//     You should have received a copy of the GNU General Public License
+//     along with this program.  If not, see <http://www.gnu.org/licenses/>.
+// 
+
+namespace KerbalEngineer.VesselSimulator
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Text;
+    using Editor;
+    using Helpers;
+    using UnityEngine;
+
+    public class EngineSim
+    {
+        private static readonly Pool<EngineSim> pool = new Pool<EngineSim>(Create, Reset);
+
+        private readonly ResourceContainer resourceConsumptions = new ResourceContainer();
+        private readonly ResourceContainer resourceFlowModes = new ResourceContainer();
+
+        public double actualThrust = 0;
+        public bool isActive = false;
+        public double isp = 0;
+        public PartSim partSim;
+        public List<AppliedForce> appliedForces = new List<AppliedForce>();
+        public float maxMach;
+
+        public double thrust = 0;
+
+        // Add thrust vector to account for directional losses
+        public Vector3 thrustVec;
+
+        private static EngineSim Create()
+        {
+            return new EngineSim();
+        }
+
+        private static void Reset(EngineSim engineSim)
+        {
+            engineSim.resourceConsumptions.Reset();
+            engineSim.resourceFlowModes.Reset();
+            engineSim.actualThrust = 0;
+            engineSim.isActive = false;
+            engineSim.isp = 0;
+            for (int i = 0; i < engineSim.appliedForces.Count; i++)
+            {
+                engineSim.appliedForces[i].Release();
+            }
+            engineSim.appliedForces.Clear();
+            engineSim.thrust = 0;
+            engineSim.maxMach = 0f;
+        }
+
+        public void Release()
+        {
+            pool.Release(this);
+        }
+
+        public static EngineSim New(PartSim theEngine,
+                         double atmosphere,
+                         float machNumber,
+                         float maxFuelFlow,
+                         float minFuelFlow,
+                         float thrustPercentage,
+                         Vector3 vecThrust,
+                         FloatCurve atmosphereCurve,
+                         bool atmChangeFlow,
+                         FloatCurve atmCurve,
+                         FloatCurve velCurve,
+                         float currentThrottle,
+                         float IspG,
+                         bool throttleLocked,
+                         List<Propellant> propellants,
+                         bool active,
+                         float resultingThrust,
+                         List<Transform> thrustTransforms,
+                        LogMsg log)
+        {
+            EngineSim engineSim = pool.Borrow();
+
+            engineSim.isp = 0.0;
+            engineSim.maxMach = 0.0f;
+            engineSim.actualThrust = 0.0;
+            engineSim.partSim = theEngine;
+            engineSim.isActive = active;
+            engineSim.thrustVec = vecThrust;
+            engineSim.resourceConsumptions.Reset();
+            engineSim.resourceFlowModes.Reset();
+            engineSim.appliedForces.Clear();
+
+            double flowRate = 0.0;
+            if (engineSim.partSim.hasVessel)
+            {
+                if (log != null) log.buf.AppendLine("hasVessel is true"); 
+
+                float flowModifier = GetFlowModifier(atmChangeFlow, atmCurve, engineSim.partSim.part.atmDensity, velCurve, machNumber, ref engineSim.maxMach);
+                engineSim.isp = atmosphereCurve.Evaluate((float)atmosphere);
+                engineSim.thrust = GetThrust(Mathf.Lerp(minFuelFlow, maxFuelFlow, GetThrustPercent(thrustPercentage)) * flowModifier, engineSim.isp);
+                engineSim.actualThrust = engineSim.isActive ? resultingThrust : 0.0;
+                if (log != null)
+                {
+                    log.buf.AppendFormat("flowMod = {0:g6}\n", flowModifier);
+                    log.buf.AppendFormat("isp     = {0:g6}\n", engineSim.isp);
+                    log.buf.AppendFormat("thrust  = {0:g6}\n", engineSim.thrust);
+                    log.buf.AppendFormat("actual  = {0:g6}\n", engineSim.actualThrust);
+                }
+
+                if (throttleLocked)
+                {
+                    if (log != null) log.buf.AppendLine("throttleLocked is true, using thrust for flowRate");
+                    flowRate = GetFlowRate(engineSim.thrust, engineSim.isp);
+                }
+                else
+                {
+                    if (currentThrottle > 0.0f && engineSim.partSim.isLanded == false)
+                    {
+                        if (log != null) log.buf.AppendLine("throttled up and not landed, using actualThrust for flowRate");
+                        flowRate = GetFlowRate(engineSim.actualThrust, engineSim.isp);
+                    }
+                    else
+                    {
+                        if (log != null) log.buf.AppendLine("throttled down or landed, using thrust for flowRate");
+                        flowRate = GetFlowRate(engineSim.thrust, engineSim.isp);
+                    }
+                }
+            }
+            else
+            {
+                if (log != null) log.buf.AppendLine("hasVessel is false");
+                float flowModifier = GetFlowModifier(atmChangeFlow, atmCurve, CelestialBodies.SelectedBody.GetDensity(BuildAdvanced.Altitude), velCurve, machNumber, ref engineSim.maxMach);
+                engineSim.isp = atmosphereCurve.Evaluate((float)atmosphere);
+                engineSim.thrust = GetThrust(Mathf.Lerp(minFuelFlow, maxFuelFlow, GetThrustPercent(thrustPercentage)) * flowModifier, engineSim.isp);
+                engineSim.actualThrust = 0d;
+                if (log != null)
+                {
+                    log.buf.AppendFormat("flowMod = {0:g6}\n", flowModifier);
+                    log.buf.AppendFormat("isp     = {0:g6}\n", engineSim.isp);
+                    log.buf.AppendFormat("thrust  = {0:g6}\n", engineSim.thrust);
+                    log.buf.AppendFormat("actual  = {0:g6}\n", engineSim.actualThrust);
+                }
+
+                if (log != null) log.buf.AppendLine("no vessel, using thrust for flowRate");
+                flowRate = GetFlowRate(engineSim.thrust, engineSim.isp);
+            }
+
+            if (log != null) log.buf.AppendFormat("flowRate = {0:g6}\n", flowRate);
+
+            float flowMass = 0f;
+            for (int i = 0; i < propellants.Count; ++i)
+            {
+                Propellant propellant = propellants[i];
+                if (!propellant.ignoreForIsp)
+                    flowMass += propellant.ratio * ResourceContainer.GetResourceDensity(propellant.id);
+            }
+
+            if (log != null) log.buf.AppendFormat("flowMass = {0:g6}\n", flowMass);
+
+            for (int i = 0; i < propellants.Count; ++i)
+            {
+                Propellant propellant = propellants[i];
+
+                if (propellant.name == "ElectricCharge" || propellant.name == "IntakeAir")
+                {
+                    continue;
+                }
+
+                double consumptionRate = propellant.ratio * flowRate / flowMass;
+                if (log != null) log.buf.AppendFormat(
+                        "Add consumption({0}, {1}:{2:d}) = {3:g6}\n",
+                        ResourceContainer.GetResourceName(propellant.id),
+                        theEngine.name,
+                        theEngine.partId,
+                        consumptionRate);
+                engineSim.resourceConsumptions.Add(propellant.id, consumptionRate);
+                engineSim.resourceFlowModes.Add(propellant.id, (double)propellant.GetFlowMode());
+            }
+
+            double thrustPerThrustTransform = engineSim.thrust / thrustTransforms.Count;
+            for (int i = 0; i < thrustTransforms.Count; i++)
+            {
+                Transform thrustTransform = thrustTransforms[i];
+                Vector3d direction = thrustTransform.forward.normalized;
+                Vector3d position = thrustTransform.position;
+
+                AppliedForce appliedForce = AppliedForce.New(direction * thrustPerThrustTransform, position);
+                engineSim.appliedForces.Add(appliedForce);
+            }
+
+            return engineSim;
+        }
+
+        public ResourceContainer ResourceConsumptions
+        {
+            get
+            {
+                return resourceConsumptions;
+            }
+        }
+
+        public static double GetExhaustVelocity(double isp)
+        {
+            return isp * Units.GRAVITY;
+        }
+
+        public static float GetFlowModifier(bool atmChangeFlow, FloatCurve atmCurve, double atmDensity, FloatCurve velCurve, float machNumber, ref float maxMach)
+        {
+            float flowModifier = 1.0f;
+            if (atmChangeFlow)
+            {
+                flowModifier = (float)(atmDensity / 1.225);
+                if (atmCurve != null)
+                {
+                    flowModifier = atmCurve.Evaluate(flowModifier);
+                }
+            }
+            if (velCurve != null)
+            {
+                flowModifier = flowModifier * velCurve.Evaluate(machNumber);
+                maxMach = velCurve.maxTime;
+            }
+            if (flowModifier < float.Epsilon)
+            {
+                flowModifier = float.Epsilon;
+            }
+            return flowModifier;
+        }
+
+        public static double GetFlowRate(double thrust, double isp)
+        {
+            return thrust / GetExhaustVelocity(isp);
+        }
+
+        public static float GetThrottlePercent(float currentThrottle, float thrustPercentage)
+        {
+            return currentThrottle * GetThrustPercent(thrustPercentage);
+        }
+
+        public static double GetThrust(double flowRate, double isp)
+        {
+            return flowRate * GetExhaustVelocity(isp);
+        }
+
+        public static float GetThrustPercent(float thrustPercentage)
+        {
+            return thrustPercentage * 0.01f;
+        }
+
+        public void DumpEngineToBuffer(StringBuilder buffer, String prefix)
+        {
+            buffer.Append(prefix);
+            buffer.AppendFormat("[thrust = {0:g6}, actual = {1:g6}, isp = {2:g6}\n", thrust, actualThrust, isp);
+        }
+
+        // A dictionary to hold a set of parts for each resource
+        Dictionary<int, HashSet<PartSim>> sourcePartSets = new Dictionary<int, HashSet<PartSim>>();
+
+        Dictionary<int, HashSet<PartSim>> stagePartSets = new Dictionary<int, HashSet<PartSim>>();
+
+        HashSet<PartSim> visited = new HashSet<PartSim>();
+
+        public void DumpSourcePartSets(String msg)
+        {
+            MonoBehaviour.print("DumpSourcePartSets " + msg);
+            foreach (int type in sourcePartSets.Keys)
+            {
+                MonoBehaviour.print("SourcePartSet for " + ResourceContainer.GetResourceName(type));
+                HashSet<PartSim> sourcePartSet = sourcePartSets[type];
+                if (sourcePartSet.Count > 0)
+                {
+                    foreach (PartSim partSim in sourcePartSet)
+                    {
+                        MonoBehaviour.print("Part " + partSim.name + ":" + partSim.partId);
+                    }
+                }
+                else
+                {
+                    MonoBehaviour.print("No parts");
+                }
+            }
+        }
+
+        public bool SetResourceDrains(List<PartSim> allParts, List<PartSim> allFuelLines, HashSet<PartSim> drainingParts)
+        {
+            LogMsg log = null;
+            //DumpSourcePartSets("before clear");
+            foreach (HashSet<PartSim> sourcePartSet in sourcePartSets.Values)
+            {
+                sourcePartSet.Clear();
+            }
+            //DumpSourcePartSets("after clear");
+
+            for (int index = 0; index < this.resourceConsumptions.Types.Count; index++)
+            {
+                int type = this.resourceConsumptions.Types[index];
+
+                HashSet<PartSim> sourcePartSet;
+                if (!sourcePartSets.TryGetValue(type, out sourcePartSet))
+                {
+                    sourcePartSet = new HashSet<PartSim>();
+                    sourcePartSets.Add(type, sourcePartSet);
+                }
+
+                switch ((ResourceFlowMode)this.resourceFlowModes[type])
+                {
+                    case ResourceFlowMode.NO_FLOW:
+                        if (partSim.resources[type] > SimManager.RESOURCE_MIN && partSim.resourceFlowStates[type] != 0)
+                        {
+                            //sourcePartSet = new HashSet<PartSim>();
+                            //MonoBehaviour.print("SetResourceDrains(" + name + ":" + partId + ") setting sources to just this");
+                            sourcePartSet.Add(partSim);
+                        }
+                        break;
+
+                    case ResourceFlowMode.ALL_VESSEL:
+                    case ResourceFlowMode.ALL_VESSEL_BALANCE:
+                        for (int i = 0; i < allParts.Count; i++)
+                        {
+                            PartSim aPartSim = allParts[i];
+                            if (aPartSim.resources[type] > SimManager.RESOURCE_MIN && aPartSim.resourceFlowStates[type] != 0)
+                            {
+                                sourcePartSet.Add(aPartSim);
+                            }
+                        }
+                        break;
+
+                    case ResourceFlowMode.STAGE_PRIORITY_FLOW:
+                    case ResourceFlowMode.STAGE_PRIORITY_FLOW_BALANCE:
+
+                        foreach (HashSet<PartSim> stagePartSet in stagePartSets.Values)
+                        {
+                            stagePartSet.Clear();
+                        }
+                        var maxStage = -1;
+
+                        //Logger.Log(type);
+                        for (int i = 0; i < allParts.Count; i++)
+                        {
+                            var aPartSim = allParts[i];
+                            if (aPartSim.resources[type] <= SimManager.RESOURCE_MIN || aPartSim.resourceFlowStates[type] == 0)
+                            {
+                                continue;
+                            }
+
+                            int stage = aPartSim.DecouplerCount();
+                            if (stage > maxStage)
+                            {
+                                maxStage = stage;
+                            }
+
+                            HashSet<PartSim> tempPartSet;
+                            if (!stagePartSets.TryGetValue(stage, out tempPartSet))
+                            {
+                                tempPartSet = new HashSet<PartSim>();
+                                stagePartSets.Add(stage, tempPartSet);
+                            }
+                            tempPartSet.Add(aPartSim);
+                        }
+
+                        for (int j = maxStage; j >= 0; j--)
+                        {
+                            HashSet<PartSim> stagePartSet;
+                            if (stagePartSets.TryGetValue(j, out stagePartSet) && stagePartSet.Count > 0)
+                            {
+                                // We have to copy the contents of the set here rather than copying the set reference or 
+                                // bad things (tm) happen
+                                foreach (PartSim aPartSim in stagePartSet)
+                                {
+                                    sourcePartSet.Add(aPartSim);
+                                }
+                                break;
+                            }
+                        }
+                        break;
+
+                    case ResourceFlowMode.STACK_PRIORITY_SEARCH:
+                        visited.Clear();
+
+                        if (SimManager.logOutput)
+                        {
+                            log = new LogMsg();
+                            log.buf.AppendLine("Find " + ResourceContainer.GetResourceName(type) + " sources for " + partSim.name + ":" + partSim.partId);
+                        }
+                        partSim.GetSourceSet(type, PhysicsGlobals.Stack_PriUsesSurf, allParts, visited, sourcePartSet, log, "");
+                        if (SimManager.logOutput && log != null)
+                        {
+                            MonoBehaviour.print(log.buf);
+                        }
+                        break;
+
+                    case ResourceFlowMode.STAGE_STACK_FLOW:
+                    case ResourceFlowMode.STAGE_STACK_FLOW_BALANCE:
+                        visited.Clear();
+
+                        if (SimManager.logOutput)
+                        {
+                            log = new LogMsg();
+                            log.buf.AppendLine("Find " + ResourceContainer.GetResourceName(type) + " sources for " + partSim.name + ":" + partSim.partId);
+                        }
+                        partSim.GetSourceSet(type, true, allParts, visited, sourcePartSet, log, "");
+                        if (SimManager.logOutput && log != null)
+                        {
+                            MonoBehaviour.print(log.buf);
+                        }
+                        break;
+
+                    default:
+                        MonoBehaviour.print("SetResourceDrains(" + partSim.name + ":" + partSim.partId + ") Unexpected flow type for " + ResourceContainer.GetResourceName(type) + ")");
+                        break;
+                }
+
+                if (SimManager.logOutput)
+                {
+                    if (sourcePartSet.Count > 0)
+                    {
+                        log = new LogMsg();
+                        log.buf.AppendLine("Source parts for " + ResourceContainer.GetResourceName(type) + ":");
+                        foreach (PartSim partSim in sourcePartSet)
+                        {
+                            log.buf.AppendLine(partSim.name + ":" + partSim.partId);
+                        }
+                        MonoBehaviour.print(log.buf);
+                    }
+                }
+
+                //DumpSourcePartSets("after " + ResourceContainer.GetResourceName(type));
+            }
+            
+            // If we don't have sources for all the needed resources then return false without setting up any drains
+            for (int i = 0; i < this.resourceConsumptions.Types.Count; i++)
+            {
+                int type = this.resourceConsumptions.Types[i];
+                HashSet<PartSim> sourcePartSet; 
+                if (!sourcePartSets.TryGetValue(type, out sourcePartSet) || sourcePartSet.Count == 0)
+                {
+                    if (SimManager.logOutput)
+                    {
+                        MonoBehaviour.print("No source of " + ResourceContainer.GetResourceName(type));
+                    }
+
+                    isActive = false;
+                    return false;
+                }
+            }
+
+            // Now we set the drains on the members of the sets and update the draining parts set
+            for (int i = 0; i < this.resourceConsumptions.Types.Count; i++)
+            {
+                int type = this.resourceConsumptions.Types[i];
+                HashSet<PartSim> sourcePartSet = sourcePartSets[type];
+                // Loop through the members of the set 
+                double amount = resourceConsumptions[type] / sourcePartSet.Count;
+                foreach (PartSim partSim in sourcePartSet)
+                {
+                    if (SimManager.logOutput)
+                    {
+                        MonoBehaviour.print(
+                            "Adding drain of " + amount + " " + ResourceContainer.GetResourceName(type) + " to " + partSim.name + ":" +
+                            partSim.partId);
+                    }
+
+                    partSim.resourceDrains.Add(type, amount);
+                    drainingParts.Add(partSim);
+                }
+            }
+            return true;
+        }
+    }
 }