epyt-flow 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

epyt_flow/simulation/scenario_simulator.py

@@ -17,15 +17,17 @@ from epyt.epanet import ToolkitConstants
 from tqdm import tqdm
 
 from .scenario_config import ScenarioConfig
-from .sensor_config import SensorConfig, SENSOR_TYPE_LINK_FLOW, SENSOR_TYPE_LINK_QUALITY, \
-    SENSOR_TYPE_NODE_DEMAND, SENSOR_TYPE_NODE_PRESSURE, SENSOR_TYPE_NODE_QUALITY, \
+from .sensor_config import SensorConfig, areaunit_to_id, massunit_to_id, qualityunit_to_id, \
+    qualityunit_to_str, MASS_UNIT_MG, \
+    SENSOR_TYPE_LINK_FLOW, SENSOR_TYPE_LINK_QUALITY, SENSOR_TYPE_NODE_DEMAND, \
+    SENSOR_TYPE_NODE_PRESSURE, SENSOR_TYPE_NODE_QUALITY, \
     SENSOR_TYPE_PUMP_STATE, SENSOR_TYPE_TANK_VOLUME, SENSOR_TYPE_VALVE_STATE, \
     SENSOR_TYPE_NODE_BULK_SPECIES, SENSOR_TYPE_LINK_BULK_SPECIES, SENSOR_TYPE_SURFACE_SPECIES
 from ..uncertainty import ModelUncertainty, SensorNoise
 from .events import SystemEvent, Leakage, ActuatorEvent, SensorFault, SensorReadingAttack, \
     SensorReadingEvent
 from .scada import ScadaData, AdvancedControlModule
-from ..topology import NetworkTopology
+from ..topology import NetworkTopology, UNITS_SIMETRIC, UNITS_USCUSTOM
 from ..utils import get_temp_folder
 
 
@@ -48,6 +50,10 @@ class ScenarioSimulator():
 
         If this is None, then 'f_inp_in' must be set with a valid path to the .inp file
         that is to be simulated.
+    epanet_verbose : `bool`, optional
+        If True, EPyT is verbose and might print messages from time to time.
+
+        The default is False.
 
     Attributes
     ----------
@@ -56,7 +62,7 @@ class ScenarioSimulator():
     """
 
     def __init__(self, f_inp_in: str = None, f_msx_in: str = None,
-                 scenario_config: ScenarioConfig = None):
+                 scenario_config: ScenarioConfig = None, epanet_verbose: bool = False):
         if f_msx_in is not None and f_inp_in is None:
             raise ValueError("'f_inp_in' must be set if 'f_msx_in' is set.")
         if f_inp_in is None and scenario_config is None:
@@ -74,6 +80,9 @@ class ScenarioSimulator():
                 raise TypeError("'scenario_config' must be an instance of " +
                                 "'epyt_flow.simulation.ScenarioConfig' but not of " +
                                 f"'{type(scenario_config)}'")
+        if not isinstance(epanet_verbose, bool):
+            raise TypeError("'epanet_verbose' must be an instance of 'bool' " +
+                            f"but not of '{type(epanet_verbose)}'")
 
         self.__f_inp_in = f_inp_in if scenario_config is None else scenario_config.f_inp_in
         self.__f_msx_in = f_msx_in if scenario_config is None else scenario_config.f_msx_in
@@ -86,37 +95,28 @@ class ScenarioSimulator():
 
         custom_epanet_lib = None
         custom_epanetmsx_lib = None
-        if sys.platform.startswith("linux"):
+        if sys.platform.startswith("linux") or sys.platform.startswith("darwin") :
             path_to_custom_libs = os.path.join(pathlib.Path(__file__).parent.resolve(),
                                                "..", "customlibs")
 
-            if os.path.isfile(os.path.join(path_to_custom_libs, "libepanet2_2.so")):
-                custom_epanet_lib = os.path.join(path_to_custom_libs, "libepanet2_2.so")
-            if os.path.isfile(os.path.join(path_to_custom_libs, "libepanetmsx2_2_0.so")):
-                custom_epanetmsx_lib = os.path.join(path_to_custom_libs, "libepanetmsx2_2_0.so")
+            libepanet_name = "libepanet2_2.so" if sys.platform.startswith("linux") \
+                else "libepanet2_2.dylib"
+            libepanetmsx_name = "libepanetmsx2_2_0.so" if sys.platform.startswith("linux") \
+                else "libepanetmsx2_2_0.dylib"
+
+            if os.path.isfile(os.path.join(path_to_custom_libs, libepanet_name)):
+                custom_epanet_lib = os.path.join(path_to_custom_libs, libepanet_name)
+            if os.path.isfile(os.path.join(path_to_custom_libs, libepanetmsx_name)):
+                custom_epanetmsx_lib = os.path.join(path_to_custom_libs, libepanetmsx_name)
 
         self.epanet_api = epanet(self.__f_inp_in, ph=self.__f_msx_in is None,
-                                 customlib=custom_epanet_lib)
+                                 customlib=custom_epanet_lib, loadfile=True,
+                                 display_msg=epanet_verbose)
 
-        bulk_species = []
-        surface_species = []
         if self.__f_msx_in is not None:
             self.epanet_api.loadMSXFile(self.__f_msx_in, customMSXlib=custom_epanetmsx_lib)
 
-            for species_id, species_type in zip(self.epanet_api.getMSXSpeciesNameID(),
-                                                self.epanet_api.getMSXSpeciesType()):
-                if species_type == "BULK":
-                    bulk_species.append(species_id)
-                elif species_type == "WALL":
-                    surface_species.append(species_id)
-
-        self.__sensor_config = SensorConfig(nodes=self.epanet_api.getNodeNameID(),
-                                            links=self.epanet_api.getLinkNameID(),
-                                            valves=self.epanet_api.getLinkValveNameID(),
-                                            pumps=self.epanet_api.getLinkPumpNameID(),
-                                            tanks=self.epanet_api.getNodeTankNameID(),
-                                            bulk_species=bulk_species,
-                                            surface_species=surface_species)
+        self.__sensor_config = self.__get_empty_sensor_config()
         if scenario_config is not None:
             if scenario_config.general_params is not None:
                 self.set_general_parameters(**scenario_config.general_params)
@@ -132,6 +132,51 @@ class ScenarioSimulator():
             for event in scenario_config.sensor_reading_events:
                 self.add_sensor_reading_event(event)
 
+    def __get_empty_sensor_config(self, node_id_to_idx: dict = None, link_id_to_idx: dict = None,
+                                  valve_id_to_idx: dict = None, pump_id_to_idx: dict = None,
+                                  tank_id_to_idx: dict = None, bulkspecies_id_to_idx: dict = None,
+                                  surfacespecies_id_to_idx: dict = None) -> SensorConfig:
+        flow_unit = self.epanet_api.api.ENgetflowunits()
+        quality_unit = qualityunit_to_id(self.epanet_api.getQualityInfo().QualityChemUnits)
+        bulk_species = []
+        surface_species = []
+        bulk_species_mass_unit = []
+        surface_species_mass_unit = []
+        surface_species_area_unit = None
+
+        if self.__f_msx_in is not None:
+            surface_species_area_unit = areaunit_to_id(self.epanet_api.getMSXAreaUnits())
+
+            for species_id, species_type, mass_unit in zip(self.epanet_api.getMSXSpeciesNameID(),
+                                                           self.epanet_api.getMSXSpeciesType(),
+                                                           self.epanet_api.getMSXSpeciesUnits()):
+                if species_type == "BULK":
+                    bulk_species.append(species_id)
+                    bulk_species_mass_unit.append(massunit_to_id(mass_unit))
+                elif species_type == "WALL":
+                    surface_species.append(species_id)
+                    surface_species_mass_unit.append(massunit_to_id(mass_unit))
+
+        return SensorConfig(nodes=self.epanet_api.getNodeNameID(),
+                            links=self.epanet_api.getLinkNameID(),
+                            valves=self.epanet_api.getLinkValveNameID(),
+                            pumps=self.epanet_api.getLinkPumpNameID(),
+                            tanks=self.epanet_api.getNodeTankNameID(),
+                            bulk_species=bulk_species,
+                            surface_species=surface_species,
+                            flow_unit=flow_unit,
+                            quality_unit=quality_unit,
+                            bulk_species_mass_unit=bulk_species_mass_unit,
+                            surface_species_mass_unit=surface_species_mass_unit,
+                            surface_species_area_unit=surface_species_area_unit,
+                            node_id_to_idx=node_id_to_idx,
+                            link_id_to_idx=link_id_to_idx,
+                            valve_id_to_idx=valve_id_to_idx,
+                            pump_id_to_idx=pump_id_to_idx,
+                            tank_id_to_idx=tank_id_to_idx,
+                            bulkspecies_id_to_idx=bulkspecies_id_to_idx,
+                            surfacespecies_id_to_idx=surfacespecies_id_to_idx)
+
     @property
     def f_inp_in(self) -> str:
         """
@@ -295,6 +340,21 @@ class ScenarioSimulator():
         return deepcopy(list(filter(lambda e: isinstance(e, SensorFault),
                                     self.__sensor_reading_events)))
 
+    @property
+    def sensor_reading_attacks(self) -> list[SensorReadingAttack]:
+        """
+        Gets all sensor reading attacks.
+
+        Returns
+        -------
+        list[:class:`~epyt_flow.simulation.events.sensor_reading_attacks.SensorReadingAttack`]
+            All sensor reading attacks.
+        """
+        self.__adapt_to_network_changes()
+
+        return deepcopy(list(filter(lambda e: isinstance(e, SensorReadingAttack)),
+                             self.__sensor_reading_events))
+
     @property
     def sensor_reading_events(self) -> list[SensorReadingEvent]:
         """
@@ -312,19 +372,6 @@ class ScenarioSimulator():
     def __adapt_to_network_changes(self):
         nodes = self.epanet_api.getNodeNameID()
         links = self.epanet_api.getLinkNameID()
-        valves = self.epanet_api.getLinkValveNameID()
-        pumps = self.epanet_api.getLinkPumpNameID()
-        tanks = self.epanet_api.getNodeTankNameID()
-        bulk_species = []
-        surface_species = []
-
-        if self.__f_msx_in is not None:
-            for species_id, species_type in zip(self.epanet_api.getMSXSpeciesNameID(),
-                                                self.epanet_api.getMSXSpeciesType()):
-                if species_type == "BULK":
-                    bulk_species.append(species_id)
-                elif species_type == "WALL":
-                    surface_species.append(species_id)
 
         node_id_to_idx = {node_id: self.epanet_api.getNodeIndex(node_id) - 1 for node_id in nodes}
         link_id_to_idx = {link_id: self.epanet_api.getLinkIndex(link_id) - 1 for link_id in links}
@@ -334,35 +381,24 @@ class ScenarioSimulator():
         bulkspecies_id_to_idx = None
         surfacespecies_id_to_idx = None
 
-        if nodes != self.__sensor_config.nodes or links != self.__sensor_config.links or \
-                valves != self.__sensor_config.valves or pumps != self.__sensor_config.pumps or \
-                tanks != self.__sensor_config.tanks or \
-                bulk_species != self.__sensor_config.bulk_species or \
-                surface_species != self.__sensor_config.surface_species:
-            # Adapt sensor configuration if anything in the network topology changed
-            new_sensor_config = SensorConfig(nodes=nodes, links=links, valves=valves, pumps=pumps,
-                                             tanks=tanks, bulk_species=bulk_species,
-                                             surface_species=surface_species,
-                                             node_id_to_idx=node_id_to_idx,
-                                             link_id_to_idx=link_id_to_idx,
-                                             valve_id_to_idx=valve_id_to_idx,
-                                             pump_id_to_idx=pump_id_to_idx,
-                                             tank_id_to_idx=tank_id_to_idx,
-                                             bulkspecies_id_to_idx=bulkspecies_id_to_idx,
-                                             surfacespecies_id_to_idx=surfacespecies_id_to_idx)
-            new_sensor_config.pressure_sensors = self.__sensor_config.pressure_sensors
-            new_sensor_config.flow_sensors = self.__sensor_config.flow_sensors
-            new_sensor_config.demand_sensors = self.__sensor_config.demand_sensors
-            new_sensor_config.quality_node_sensors = self.__sensor_config.quality_node_sensors
-            new_sensor_config.quality_link_sensors = self.__sensor_config.quality_link_sensors
-            new_sensor_config.pump_state_sensors = self.__sensor_config.pump_state_sensors
-            new_sensor_config.valve_state_sensors = self.__sensor_config.valve_state_sensors
-            new_sensor_config.tank_volume_sensors = self.__sensor_config.tank_volume_sensors
-            new_sensor_config.bulk_species_node_sensors = self.__sensor_config.bulk_species_node_sensors
-            new_sensor_config.bulk_species_link_sensors = self.__sensor_config.bulk_species_link_sensors
-            new_sensor_config.surface_species_sensors = self.__sensor_config.surface_species_sensors
-
-            self.__sensor_config = new_sensor_config
+        # Adapt sensor configuration to potential cahnges in the network's topology
+        new_sensor_config = self.__get_empty_sensor_config(node_id_to_idx, link_id_to_idx,
+                                                           valve_id_to_idx, pump_id_to_idx,
+                                                           tank_id_to_idx, bulkspecies_id_to_idx,
+                                                           surfacespecies_id_to_idx)
+        new_sensor_config.pressure_sensors = self.__sensor_config.pressure_sensors
+        new_sensor_config.flow_sensors = self.__sensor_config.flow_sensors
+        new_sensor_config.demand_sensors = self.__sensor_config.demand_sensors
+        new_sensor_config.quality_node_sensors = self.__sensor_config.quality_node_sensors
+        new_sensor_config.quality_link_sensors = self.__sensor_config.quality_link_sensors
+        new_sensor_config.pump_state_sensors = self.__sensor_config.pump_state_sensors
+        new_sensor_config.valve_state_sensors = self.__sensor_config.valve_state_sensors
+        new_sensor_config.tank_volume_sensors = self.__sensor_config.tank_volume_sensors
+        new_sensor_config.bulk_species_node_sensors = self.__sensor_config.bulk_species_node_sensors
+        new_sensor_config.bulk_species_link_sensors = self.__sensor_config.bulk_species_link_sensors
+        new_sensor_config.surface_species_sensors = self.__sensor_config.surface_species_sensors
+
+        self.__sensor_config = new_sensor_config
 
     def close(self):
         """
@@ -381,22 +417,185 @@ class ScenarioSimulator():
     def __exit__(self, *args):
         self.close()
 
+    def save_to_epanet_file(self, inp_file_path: str, msx_file_path: str = None,
+                            export_sensor_config: bool = True) -> None:
+        """
+        Exports this scenario to EPANET files -- i.e. an .inp file
+        and (optionally) a .msx file if EPANET-MSX was loaded.
+
+        Parameters
+        ----------
+        inp_file_path : `str`
+            Path to the .inp file where this scenario will be stored.
+
+            If 'inp_file_path' is None, 'msx_file_path' must not be None!
+        msx_file_path : `str`, optional
+            Path to the .msx file where this MSX component of this scneario will be stored.
+
+            Note that this is only applicable if EPANET-MSX was loaded.
+
+            The default is None.
+        export_sensor_config : `bool`, optional
+            If True, the current sensor placement is exported as well.
+
+            The default is True.
+        """
+        if inp_file_path is None and msx_file_path is None:
+            raise ValueError("At least one of the paths (.inp and .msx) must not be None")
+        if inp_file_path is not None:
+            if not isinstance(inp_file_path, str):
+                raise TypeError("'inp_file_path' must be an instance of 'str' " +
+                                f"but not of '{type(inp_file_path)}'")
+        if msx_file_path is not None:
+            if not isinstance(msx_file_path, str):
+                raise TypeError("msx_file_path' msut be an instance of 'str' " +
+                                f"but not of {type(msx_file_path)}")
+        if not isinstance(export_sensor_config, bool):
+            raise TypeError("'export_sensor_config' must be an instance of 'bool' " +
+                            f"but not of '{type(export_sensor_config)}'")
+
+        def __override_report_section(file_in: str, report_desc: str) -> None:
+            with open(file_in, mode="r+", encoding="utf-8") as f_in:
+                # Find and remove exiting REPORT section
+                content = f_in.read()
+                try:
+                    report_section_start_idx = content.index("[REPORT]")
+                    report_section_end_idx = content.index("[", report_section_start_idx + 1)
+
+                    content = content[:report_section_start_idx] + content[report_section_end_idx:]
+                    f_in.seek(0)
+                    f_in.write(content)
+                    f_in.truncate()
+                except ValueError:
+                    pass
+
+                # Write new REPORT section in the very end of the file
+                write_end_section = False
+                try:
+                    end_idx = content.index("[END]")
+                    write_end_section = True
+                    f_in.seek(end_idx)
+                except ValueError:
+                    pass
+                f_in.write(report_desc)
+                if write_end_section is True:
+                    f_in.write("\n[END]")
+
+        if inp_file_path is not None:
+            self.epanet_api.saveInputFile(inp_file_path)
+
+            if export_sensor_config is True:
+                report_desc = "\n\n[REPORT]\n"
+                report_desc += "ENERGY YES\n"
+                report_desc += "STATUS YES\n"
+
+                nodes = []
+                links = []
+
+                # Parse sensor config
+                pressure_sensors = self.__sensor_config.pressure_sensors
+                if len(pressure_sensors) != 0:
+                    report_desc += "Pressure YES\n"
+                    nodes += pressure_sensors
+
+                flow_sensors = self.__sensor_config.flow_sensors
+                if len(flow_sensors) != 0:
+                    report_desc += "Flow YES\n"
+                    links += flow_sensors
+
+                demand_sensors = self.__sensor_config.demand_sensors
+                if len(demand_sensors) != 0:
+                    report_desc += "Demand YES\n"
+                    nodes += demand_sensors
+
+                node_quality_sensors = self.__sensor_config.quality_node_sensors
+                if len(node_quality_sensors) != 0:
+                    report_desc += "Quality YES\n"
+                    nodes += node_quality_sensors
+
+                link_quality_sensors = self.__sensor_config.quality_link_sensors
+                if len(link_quality_sensors) != 0:
+                    if len(node_quality_sensors) == 0:
+                        report_desc += "Quality YES\n"
+                    links += link_quality_sensors
+
+                # Create final REPORT section
+                nodes = list(set(nodes))
+                links = list(set(links))
+
+                if len(nodes) != 0:
+                    if set(nodes) == set(self.__sensor_config.nodes):
+                        nodes = ["ALL"]
+                    report_desc += f"NODES {' '.join(nodes)}\n"
+
+                if len(links) != 0:
+                    if set(links) == set(self.__sensor_config.links):
+                        links = ["ALL"]
+                    report_desc += f"LINKS {' '.join(links)}\n"
+
+                __override_report_section(inp_file_path, report_desc)
+
+        if self.__f_msx_in is not None and msx_file_path is not None:
+            self.epanet_api.saveMSXFile(msx_file_path)
+
+            if export_sensor_config is True:
+                report_desc = "\n\n[REPORT]\n"
+                species = []
+                nodes = []
+                links = []
+
+                # Parse sensor config
+                bulk_species_node_sensors = self.__sensor_config.bulk_species_node_sensors
+                for bulk_species_id in bulk_species_node_sensors.keys():
+                    species.append(bulk_species_id)
+                    nodes += bulk_species_node_sensors[bulk_species_id]
+
+                bulk_species_link_sensors = self.__sensor_config.bulk_species_link_sensors
+                for bulk_species_id in bulk_species_link_sensors.keys():
+                    species.append(bulk_species_id)
+                    links += bulk_species_link_sensors[bulk_species_id]
+
+                surface_species_link_sensors = self.__sensor_config.surface_species_sensors
+                for surface_species_id in surface_species_link_sensors.keys():
+                    species.append(surface_species_id)
+                    links += surface_species_link_sensors[surface_species_id]
+
+                nodes = list(set(nodes))
+                links = list((set(links)))
+                species = list(set(species))
+
+                # Create REPORT section
+                if len(nodes) != 0:
+                    if set(nodes) == set(self.__sensor_config.nodes):
+                        nodes = ["ALL"]
+                    report_desc += f"NODES {' '.join(nodes)}\n"
+
+                if len(links) != 0:
+                    if set(links) == set(self.__sensor_config.links):
+                        links = ["ALL"]
+                    report_desc += f"LINKS {' '.join(links)}\n"
+
+                for species_id in species:
+                    report_desc += f"SPECIES {species_id} YES\n"
+
+                __override_report_section(msx_file_path, report_desc)
+
     def get_flow_units(self) -> int:
         """
         Gets the flow units.
 
         Will be one of the following EPANET toolkit constants:
 
-            - EN_CFS = 0
-            - EN_GPM = 1
-            - EN_MGD = 2
-            - EN_IMGD = 3
-            - EN_AFD = 4
-            - EN_LPS = 5
-            - EN_LPM = 6
-            - EN_MLD = 7
-            - EN_CMH = 8
-            - EN_CMD = 9
+            - EN_CFS = 0 (cu foot/sec)
+            - EN_GPM = 1 (gal/min)
+            - EN_MGD = 2 (Million gal/day)
+            - EN_IMGD = 3 (Imperial MGD)
+            - EN_AFD = 4 (ac-foot/day)
+            - EN_LPS = 5 (liter/sec)
+            - EN_LPM = 6 (liter/min)
+            - EN_MLD = 7 (Megaliter/day)
+            - EN_CMH = 8 (cubic meter/hr)
+            - EN_CMD = 9 (cubic meter/day)
 
         Returns
         -------
@@ -405,6 +604,27 @@ class ScenarioSimulator():
         """
         return self.epanet_api.api.ENgetflowunits()
 
+    def get_units_category(self) -> int:
+        """
+        Gets the category of units -- i.e. US Customary or SI Metric units.
+
+        Will be one of the following constants:
+
+            - UNITS_USCUSTOM = 0  (US Customary)
+            - UNITS_SIMETRIC = 1  (SI Metric)
+
+        Returns
+        -------
+        `int`
+            Units category.
+        """
+        if self.get_flow_units() in [ToolkitConstants.EN_CFS, ToolkitConstants.EN_GPM,
+                                     ToolkitConstants.EN_MGD, ToolkitConstants.EN_IMGD,
+                                     ToolkitConstants.EN_AFD]:
+            return UNITS_USCUSTOM
+        else:
+            return UNITS_SIMETRIC
+
     def get_hydraulic_time_step(self) -> int:
         """
         Gets the hydraulic time step -- i.e. time step in the hydraulic simulation.
@@ -471,7 +691,7 @@ class ScenarioSimulator():
         return {"code": qual_info.QualityCode,
                 "type": qual_info.QualityType,
                 "chemical_name": qual_info.QualityChemName,
-                "units": qual_info.QualityChemUnits,
+                "units": qualityunit_to_id(qual_info.QualityChemUnits),
                 "trace_node_id": qual_info.TraceNode}
 
     def get_scenario_config(self) -> ScenarioConfig:
@@ -518,6 +738,10 @@ class ScenarioSimulator():
         n_quantities = self.epanet_api.getNodeCount() * 3 + self.epanet_api.getNodeTankCount() + \
                        self.epanet_api.getLinkValveCount() + self.epanet_api.getLinkPumpCount() + \
                        self.epanet_api.getLinkCount() * 2
+
+        if self.__f_msx_in is not None:
+            n_quantities += self.epanet_api.getLinkCount() * 2 + self.epanet_api.getNodeCount()
+
         n_bytes_per_quantity = 64
 
         return n_time_steps * n_quantities * n_bytes_per_quantity * .000001
@@ -537,6 +761,9 @@ class ScenarioSimulator():
         nodes_id = self.epanet_api.getNodeNameID()
         nodes_elevation = self.epanet_api.getNodeElevations()
         nodes_type = [self.epanet_api.TYPENODE[i] for i in self.epanet_api.getNodeTypeIndex()]
+        nodes_coord = [self.epanet_api.api.ENgetcoord(node_idx)
+                       for node_idx in self.epanet_api.getNodeIndex()]
+        node_tank_names = self.epanet_api.getNodeTankNameID()
 
         links_id = self.epanet_api.getLinkNameID()
         links_data = self.epanet_api.getNodesConnectingLinksID()
@@ -549,8 +776,16 @@ class ScenarioSimulator():
 
         # Build graph describing the topology
         nodes = []
-        for node, node_elevation, node_type in zip(nodes_id, nodes_elevation, nodes_type):
-            nodes.append((node, {"elevation": node_elevation, "type": node_type}))
+        for node_id, node_elevation, node_type, node_coord in zip(nodes_id, nodes_elevation,
+                                                                  nodes_type, nodes_coord):
+            node_info = {"elevation": node_elevation,
+                         "coord": node_coord,
+                         "type": node_type}
+            if node_type == "TANK":
+                node_tank_idx = node_tank_names.index(node_id) + 1
+                node_info["diameter"] = float(self.epanet_api.getNodeTankDiameter(node_tank_idx))
+
+            nodes.append((node_id, node_info))
 
         links = []
         for link_id, link, diameter, length, roughness_coeff, bulk_coeff, wall_coeff, loss_coeff \
@@ -561,7 +796,8 @@ class ScenarioSimulator():
                                           "bulk_coeff": bulk_coeff, "wall_coeff": wall_coeff,
                                           "loss_coeff": loss_coeff}))
 
-        return NetworkTopology(f_inp=self.f_inp_in, nodes=nodes, links=links)
+        return NetworkTopology(f_inp=self.f_inp_in, nodes=nodes, links=links,
+                               units=self.get_units_category())
 
     def randomize_demands(self) -> None:
         """
@@ -974,6 +1210,9 @@ class ScenarioSimulator():
         :class:`~epyt_flow.simulation.scada.scada_data.ScadaData`
             Quality simulation results as SCADA data.
         """
+        if self.__f_msx_in is None:
+            raise ValueError("No .msx file specified")
+
         result = None
 
         gen = self.run_advanced_quality_simulation_as_generator
@@ -1035,6 +1274,9 @@ class ScenarioSimulator():
             Generator containing the current EPANET-MSX simulation results as SCADA data
             (i.e. species concentrations).
         """
+        if self.__f_msx_in is None:
+            raise ValueError("No .msx file specified")
+
         # Load pre-computed hydraulics
         self.epanet_api.useMSXHydraulicFile(hyd_file_in)
 
@@ -1508,8 +1750,7 @@ class ScenarioSimulator():
                 # Fetch data
                 pressure_data = self.epanet_api.getNodePressure().reshape(1, -1)
                 flow_data = self.epanet_api.getLinkFlows().reshape(1, -1)
-                demand_data = self.epanet_api.getNodeActualDemand().reshape(1,
-                                                                            -1)  # TODO: Does not go back after emitter coefficient is changed back to zero
+                demand_data = self.epanet_api.getNodeActualDemand().reshape(1, -1)
                 quality_node_data = self.epanet_api.getNodeActualQuality().reshape(1, -1)
                 quality_link_data = self.epanet_api.getLinkActualQuality().reshape(1, -1)
                 pumps_state_data = self.epanet_api.getLinkPumpState().reshape(1, -1)
@@ -1610,7 +1851,7 @@ class ScenarioSimulator():
     def set_general_parameters(self, demand_model: dict = None, simulation_duration: int = None,
                                hydraulic_time_step: int = None, quality_time_step: int = None,
                                reporting_time_step: int = None, reporting_time_start: int = None,
-                               flow_units: int = None, quality_model: dict = None) -> None:
+                               flow_units_id: int = None, quality_model: dict = None) -> None:
         """
         Sets some general parameters.
 
@@ -1655,22 +1896,22 @@ class ScenarioSimulator():
             Start time (in seconds) at which reporting of hydraulic and quality states starts.
 
             The default is None.
-        flow_units : `int`, optional
+        flow_units_id : `int`, optional
             Specifies the flow units -- i.e. all flows will be reported in these units.
             If None, the units from the .inp file will be used.
 
             Must be one of the following EPANET toolkit constants:
 
-                - EN_CFS = 0
-                - EN_GPM = 1
-                - EN_MGD = 2
-                - EN_IMGD = 3
-                - EN_AFD = 4
-                - EN_LPS = 5
-                - EN_LPM = 6
-                - EN_MLD = 7
-                - EN_CMH = 8
-                - EN_CMD = 9
+                - EN_CFS  = 0  (cubic foot/sec)
+                - EN_GPM  = 1  (gal/min)
+                - EN_MGD  = 2  (Million gal/day)
+                - EN_IMGD = 3  (Imperial MGD)
+                - EN_AFD  = 4  (ac-foot/day)
+                - EN_LPS  = 5  (liter/sec)
+                - EN_LPM  = 6  (liter/min)
+                - EN_MLD  = 7  (Megaliter/day)
+                - EN_CMH  = 8  (cubic meter/hr)
+                - EN_CMD  = 9  (cubic meter/day)
 
             The default is None.
         quality_model : `dict`, optional
@@ -1691,8 +1932,7 @@ class ScenarioSimulator():
             if not isinstance(simulation_duration, int) or simulation_duration <= 0:
                 raise ValueError("'simulation_duration' must be a positive integer specifying " +
                                  "the number of seconds to simulate")
-            self.epanet_api.setTimeSimulationDuration(simulation_duration)  # TODO: Changing the simulation
-            # duration from .inp file seems to break EPANET-MSX
+            self.epanet_api.setTimeSimulationDuration(simulation_duration)
 
         if hydraulic_time_step is not None:
             if not isinstance(hydraulic_time_step, int) or hydraulic_time_step <= 0:
@@ -1727,29 +1967,29 @@ class ScenarioSimulator():
                                  "greater than the hydraulic time step")
             self.epanet_api.setTimeQualityStep(quality_time_step)
 
-        if flow_units is not None:
-            if flow_units == ToolkitConstants.EN_CFS:
+        if flow_units_id is not None:
+            if flow_units_id == ToolkitConstants.EN_CFS:
                 self.epanet_api.setFlowUnitsCFS()
-            elif flow_units == ToolkitConstants.EN_GPM:
+            elif flow_units_id == ToolkitConstants.EN_GPM:
                 self.epanet_api.setFlowUnitsGPM()
-            elif flow_units == ToolkitConstants.EN_MGD:
+            elif flow_units_id == ToolkitConstants.EN_MGD:
                 self.epanet_api.setFlowUnitsMGD()
-            elif flow_units == ToolkitConstants.EN_IMGD:
+            elif flow_units_id == ToolkitConstants.EN_IMGD:
                 self.epanet_api.setFlowUnitsIMGD()
-            elif flow_units == ToolkitConstants.EN_AFD:
+            elif flow_units_id == ToolkitConstants.EN_AFD:
                 self.epanet_api.setFlowUnitsAFD()
-            elif flow_units == ToolkitConstants.EN_LPS:
+            elif flow_units_id == ToolkitConstants.EN_LPS:
                 self.epanet_api.setFlowUnitsLPS()
-            elif flow_units == ToolkitConstants.EN_LPM:
+            elif flow_units_id == ToolkitConstants.EN_LPM:
                 self.epanet_api.setFlowUnitsLPM()
-            elif flow_units == ToolkitConstants.EN_MLD:
+            elif flow_units_id == ToolkitConstants.EN_MLD:
                 self.epanet_api.setFlowUnitsMLD()
-            elif flow_units == ToolkitConstants.EN_CMH:
+            elif flow_units_id == ToolkitConstants.EN_CMH:
                 self.epanet_api.setFlowUnitsCMH()
-            elif flow_units == ToolkitConstants.EN_CMD:
+            elif flow_units_id == ToolkitConstants.EN_CMD:
                 self.epanet_api.setFlowUnitsCMD()
             else:
-                raise ValueError(f"Unknown flow units '{flow_units}'")
+                raise ValueError(f"Unknown flow units '{flow_units_id}'")
 
         if quality_model is not None:
             if quality_model["type"] == "NONE":
@@ -1758,7 +1998,7 @@ class ScenarioSimulator():
                 self.epanet_api.setQualityType("age")
             elif quality_model["type"] == "CHEM":
                 self.epanet_api.setQualityType("chem", quality_model["chemical_name"],
-                                               quality_model["units"])
+                                               qualityunit_to_str(quality_model["units"]))
             elif quality_model["type"] == "TRACE":
                 self.epanet_api.setQualityType("trace", quality_model["trace_node_id"])
             else:
@@ -1809,7 +2049,7 @@ class ScenarioSimulator():
         self.set_general_parameters(quality_model={"type": "AGE"})
 
     def enable_chemical_analysis(self, chemical_name: str = "Chlorine",
-                                 chemical_units: str = "mg/L") -> None:
+                                 chemical_units: int = MASS_UNIT_MG) -> None:
         """
         Sets chemical analysis.
 
@@ -1823,9 +2063,13 @@ class ScenarioSimulator():
             The default is "Chlorine".
         chemical_units : `str`, optional
             Units that the chemical is measured in.
-            Either "mg/L" or "ug/L".
 
-            The default is "mg/L".
+            Must be one of the following constants:
+
+                - MASS_UNIT_MG = 4  (mg/L)
+                - MASS_UNIT_UG = 5  (ug/L)
+
+            The default is MASS_UNIT_MG.
         """
         self.__adapt_to_network_changes()