Plot units (#296)

K20shores · boulderdaze · web-flow · commit 1c09eca93bae · 2024-12-20T13:26:51.000-06:00
* adding an option to specify plot units

* updating readme

* removing music box output

* correcting test

* Update src/acom_music_box/utils.py

Co-authored-by: Jiwon Gim &lt;55209567+boulderdaze@users.noreply.github.com&gt;

---------

Co-authored-by: Jiwon Gim &lt;55209567+boulderdaze@users.noreply.github.com&gt;
diff --git a/README.md b/README.md
@@ -86,6 +86,18 @@ music_box -e TS1 --plot O3 --plot PAN,HF
 
 Note that the windows may overlap each other
 
+By default all plot units are in `mol m-3`. You can see a list of unit options to specify with `--plot-output-unit`
+
+```
+music_box -h
+```
+
+It is used like this
+
+```
+ music_box -e TS1 --output-format csv --plot O3 --plot-output-unit "ppb"
+```
+
 ### gnuplot
 If you want ascii plots (maybe you're running over ssh and can't view a graphical window), you can set
 the plot tool to gnuplo (`--plot-tool gnuplot`) to view some output
diff --git a/src/acom_music_box/constants.py b/src/acom_music_box/constants.py
@@ -1,3 +1,3 @@
-BOLTZMANN_CONSTANT = 1.380649e-23       # joules / Kelvin
-AVOGADRO_CONSTANT = 6.02214076e23       # / mole
-GAS_CONSTANT = BOLTZMANN_CONSTANT * AVOGADRO_CONSTANT   # joules / Kelvin-mole
+BOLTZMANN_CONSTANT = 1.380649e-23       # joules/Kelvin (kg m2 s-2 K-1)
+AVOGADRO_CONSTANT = 6.02214076e23       # mol-1
+GAS_CONSTANT = BOLTZMANN_CONSTANT * AVOGADRO_CONSTANT   # joules / Kelvin-mole (kg m2 s-2 K-1 mol-1)
diff --git a/src/acom_music_box/main.py b/src/acom_music_box/main.py
@@ -3,8 +3,8 @@
 import datetime
 import logging
 import os
-import sys
 from acom_music_box import MusicBox, Examples, __version__, DataOutput, PlotOutput
+from acom_music_box.utils import get_available_units
 
 
 def format_examples_help(examples):
@@ -67,6 +67,13 @@ def parse_arguments():
         default='matplotlib',
         help='Choose plotting tool: gnuplot or matplotlib (default: matplotlib).'
     )
+    parser.add_argument(
+        '--plot-output-unit',
+        type=str,
+        choices=get_available_units(),
+        default='mol m-3',
+        help='Specify the output unit for plotting concentrations.'
+    )
     return parser.parse_args()
 
 
diff --git a/src/acom_music_box/plot_output.py b/src/acom_music_box/plot_output.py
@@ -4,6 +4,7 @@
 import subprocess
 import os
 import tempfile
+from acom_music_box.utils import convert_from_number_density  # Assuming a utility function for unit conversion
 
 logger = logging.getLogger(__name__)
 
@@ -54,6 +55,7 @@ def __init__(self, df, args):
 
         self.df = df.copy(deep=True)
         self.args = args
+        self.output_unit = args.plot_output_unit if args.plot_output_unit else 'mol m-3'
         if self.args.plot:
             self.species_list = [self._format_species_list(group.split(',')) for group in self.args.plot]
         else:
@@ -88,6 +90,30 @@ def _format_species_list(self, species_list):
 
         return plot_list
 
+    def _convert_units(self, data):
+        """
+        Convert the data to the specified output unit.
+
+        Parameters
+        ----------
+        data : pandas.DataFrame
+            The DataFrame containing the data to be converted.
+
+        Returns
+        -------
+        pandas.DataFrame
+            The DataFrame with data converted to the specified unit.
+        """
+        converted_data = data.copy()
+        temperature = data['ENV.temperature']
+        pressure = data['ENV.pressure']
+        for column in data.columns:
+            if ('time' in column) or ('ENV' in column):
+                continue
+            converted_data[column] = convert_from_number_density(data[column], self.output_unit, temperature=temperature, pressure=pressure)  # Assuming standard conditions
+            converted_data.rename(columns={column: column.replace('mol m-3', self.output_unit)}, inplace=True)
+        return converted_data
+
     def _plot_with_gnuplot(self):
         """
         Plot the specified species using gnuplot.
@@ -114,8 +140,8 @@ def _plot_with_gnuplot(self):
                     gnuplot_command = f"""
                 set datafile separator ",";
                 set terminal dumb size 120,25;
-                set xlabel 'Time';
-                set ylabel 'Value';
+                set xlabel 'Time [s]';
+                set ylabel 'Concentration [{self.output_unit}]';
                 set title 'Time vs Species';
                 plot {plot_commands}
                 """
@@ -139,7 +165,7 @@ def _plot_with_matplotlib(self):
             fig, ax = plt.subplots()
             indexed[species_group].plot(ax=ax)
 
-            ax.set(xlabel='Time [s]', ylabel='Concentration [mol m-3]', title='Time vs Species')
+            ax.set(xlabel='Time [s]', ylabel=f'Concentration [{self.output_unit}]', title='Time vs Species')
 
             ax.spines[:].set_visible(False)
             ax.spines['left'].set_visible(True)
@@ -167,6 +193,7 @@ def plot(self):
             logger.debug("No species provided for plotting.")
             return
 
+        self.df = self._convert_units(self.df)
         if self.args.plot_tool == 'gnuplot':
             self._plot_with_gnuplot()
         else:
diff --git a/src/acom_music_box/utils.py b/src/acom_music_box/utils.py
@@ -1,6 +1,23 @@
 import re
 from .constants import GAS_CONSTANT, AVOGADRO_CONSTANT
-
+import numpy as np
+
+# The possible units we can convert to and from
+# functions that do conversions update this dictionary for their units in 
+# the appropriate way
+unit_conversions = {
+    'mol m-3': 0,
+    'mol cm-3': 0,
+    'molec m-3': 0,
+    'molecule m-3': 0,
+    'molec cm-3': 0,
+    'molecule cm-3': 0,
+    'ppth': 0,
+    'ppm': 0,
+    'ppb': 0,
+    'ppt': 0,
+    'mol mol-1': 0
+}
 
 def extract_unit(data, key):
     """Extract the value and unit from the key in data."""
@@ -87,6 +104,7 @@ def convert_temperature(data, key):
 def convert_concentration(data, key, temperature, pressure):
     """
     Convert the concentration from the input data to moles per cubic meter.
+    This function assumes you are passing data from a music box configuration.
 
     Args:
         data (dict): The input data.
@@ -97,36 +115,112 @@ def convert_concentration(data, key, temperature, pressure):
         float: The concentration in moles per cubic meter.
     """
     concentration_value, unit = extract_unit(data, key)
+    return convert_to_number_density(concentration_value, unit, temperature, pressure)
+
+
+def convert_to_number_density(data, input_unit, temperature, pressure):
+    """
+    Convert from some other units to mol m-3
+
+    Args:
+        data (float): The data to convert in the input unit.
+        input_unit (str): The input units
+        temperature (float): The temperature in Kelvin.
+        pressure (float): The pressure in Pascals.
+    Returns:
+        float: The data in the output unit.
+    """
+
     air_density = calculate_air_density(temperature, pressure)
 
-    unit_conversions = {
+    conversions = {a: b for a, b in unit_conversions.items()}
+    conversions.update({
         'mol m-3': 1,  # mol m-3 is the base unit
         'mol cm-3': 1e6,  # cm3 m-3
         'molec m-3': 1 / AVOGADRO_CONSTANT,  # mol
         'molecule m-3': 1 / AVOGADRO_CONSTANT,  # mol
         'molec cm-3': 1e6 / AVOGADRO_CONSTANT,  # mol cm3 m-3
         'molecule cm-3': 1e6 / AVOGADRO_CONSTANT,  # mol cm3 m-3
-        'ppth': 1e-3 * air_density,  # moles / m^3
-        'ppm': 1e-6 * air_density,  # moles / m^3
-        'ppb': 1e-9 * air_density,  # moles / m^3
-        'ppt': 1e-12 * air_density,  # moles / m^3
-        'mol mol-1': 1 * air_density  # moles / m^3
-    }
-
-    if unit in unit_conversions:
-        return concentration_value * unit_conversions[unit]
+        'ppth': 1e-3 * air_density,  # m3 mol-1
+        'ppm': 1e-6 * air_density,  # m3 mol-1
+        'ppb': 1e-9 * air_density,  # m3 mol-1
+        'ppt': 1e-12 * air_density,  # m3 mol-1
+        'mol mol-1': 1 * air_density  # m3 mol-1
+    })
+
+    if input_unit not in conversions:
+        raise ValueError(f"Unable to convert from {input_unit} to mol m-3")
+
+    conversion_factor = conversions.get(input_unit)
+
+    if isinstance(data, np.ndarray):
+        return data * conversion_factor
+    elif isinstance(data, list):
+        return [x * conversion_factor for x in data]
     else:
-        raise ValueError(f"Unsupported concentration unit: {unit}")
+        return data * conversion_factor
+
+
+def convert_from_number_density(data, output_unit, temperature, pressure):
+    """
+    Convert from mol m-3 to some other units
+
+    Args:
+        data (float): The data to convert in mol m-3.
+        output_unit (str): The output units
+        temperature (float): The temperature in Kelvin.
+        pressure (float): The pressure in Pascals.
+    Returns:
+        float: The data in the output unit.
+    """
 
+    air_density = calculate_air_density(temperature, pressure)
+
+    conversions = {a: b for a, b in unit_conversions.items()}
+    conversions.update({
+        'mol m-3': 1,  # mol m-3 is the base unit
+        'mol cm-3': 1e-6,  # m3 cm-3
+        'molec m-3': 1 * AVOGADRO_CONSTANT,  # mol-1
+        'molecule m-3': 1 * AVOGADRO_CONSTANT,  # mol-1
+        'molec cm-3': 1e-6 * AVOGADRO_CONSTANT,  # m3 cm-3 mol-1
+        'molecule cm-3': 1e-6 * AVOGADRO_CONSTANT,  # m3 cm-3 mol-1
+        'ppth': 1e3 / air_density,  # unitless
+        'ppm': 1e6 / air_density,  # unitless
+        'ppb': 1e9 / air_density,  # unitless
+        'ppt': 1e12 / air_density,  # unitless
+        'mol mol-1': 1 / air_density  # unitless
+    })
+
+    if output_unit not in conversions:
+        raise ValueError(f"Unable to convert from mol m-3 to {output_unit}")
+
+    conversion_factor = conversions.get(output_unit)
+
+    if isinstance(data, np.ndarray):
+        return data * conversion_factor
+    elif isinstance(data, list):
+        return [x * conversion_factor for x in data]
+    else:
+        return data * conversion_factor
 
 def calculate_air_density(temperature, pressure):
     """
-    Calculate the air density in moles/m^3.
+    Calculate the air density in moles m-3
 
     Args:
         temperature (float): The temperature in Kelvin.
         pressure (float): The pressure in Pascals.
     Returns:
-        float: The air density in moles/m^3.
+        float: The air density in moles m-3
     """
     return pressure / (GAS_CONSTANT * temperature)
+
+
+def get_available_units():
+    """
+    Get the list of available units for conversion.
+
+    Returns:
+        list: The list of available units.
+    """
+    return list(unit_conversions.keys())
diff --git a/tests/unit/test_plot_output.py b/tests/unit/test_plot_output.py
@@ -4,7 +4,6 @@
 import shutil
 from argparse import Namespace
 import matplotlib
-import subprocess
 
 matplotlib.use('Agg')  # Use a non-interactive backend
 
@@ -17,21 +16,23 @@ def setUp(self):
             'time': [0, 1, 2],
             'CONC.A': [1, 2, 3],
             'CONC.B': [4, 5, 6],
-            'CONC.C': [7, 8, 9]
+            'CONC.C': [7, 8, 9],
+            'ENV.temperature': [298.15, 298.15, 298.15],
+            'ENV.pressure': [101325, 101325, 101325]
         })
 
     def test_format_species_list(self):
-        args = Namespace(plot=['A', 'B'], plot_tool='matplotlib')
+        args = Namespace(plot=['A', 'B'], plot_tool='matplotlib', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         expected_list = [['CONC.A'], ['CONC.B']]
         self.assertEqual(plot_output.species_list, expected_list)
 
-        args = Namespace(plot=['CONC.A', 'CONC.B'], plot_tool='matplotlib')
+        args = Namespace(plot=['CONC.A', 'CONC.B'], plot_tool='matplotlib', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         self.assertEqual(plot_output.species_list, expected_list)
 
     def test_plot_with_gnuplot(self):
-        args = Namespace(plot=['A', 'B'], plot_tool='gnuplot')
+        args = Namespace(plot=['A', 'B'], plot_tool='gnuplot', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         if shutil.which('gnuplot') is None:
             with self.assertRaises(FileNotFoundError):
@@ -40,12 +41,12 @@ def test_plot_with_gnuplot(self):
             plot_output.plot()
 
     def test_plot_with_matplotlib(self):
-        args = Namespace(plot=['A', 'B'], plot_tool='matplotlib')
+        args = Namespace(plot=['A', 'B'], plot_tool='matplotlib', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         plot_output.plot()
 
     def test_multiple_groups_with_gnuplot(self):
-        args = Namespace(plot=['A,B', 'C'], plot_tool='gnuplot')
+        args = Namespace(plot=['A,B', 'C'], plot_tool='gnuplot', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         if shutil.which('gnuplot') is None:
             with self.assertRaises(FileNotFoundError):
@@ -54,7 +55,7 @@ def test_multiple_groups_with_gnuplot(self):
             plot_output.plot()
 
     def test_multiple_groups_with_matplotlib(self):
-        args = Namespace(plot=['A,B', 'C'], plot_tool='matplotlib')
+        args = Namespace(plot=['A,B', 'C'], plot_tool='matplotlib', plot_output_unit='mol m-3')
         plot_output = PlotOutput(self.df, args)
         plot_output.plot()
 
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -4,7 +4,8 @@
     convert_pressure,
     convert_temperature,
     convert_concentration,
-    calculate_air_density
+    calculate_air_density,
+    convert_from_number_density
 )
 import math
 
@@ -60,3 +61,20 @@ def test_invalid_concentration():
     data = {'invalid_concentration': 100}
     with pytest.raises(ValueError):
         convert_concentration(data, 'invalid_concentration', 298.15, 101325)
+
+@pytest.mark.parametrize("data, output_unit, temperature, pressure, expected",
+[
+    (1, 'mol m-3', 298.15, 101325, 1),
+    (1, 'mol cm-3', 298.15, 101325, 1e-6),
+    (1, 'molec m-3', 298.15, 101325, 1 * 6.02214076e+23),
+    (1, 'molec cm-3', 298.15, 101325, 1e-6 * 6.02214076e+23),
+    (1, 'molecule m-3', 298.15, 101325, 1 * 6.02214076e+23),
+    (1, 'molecule cm-3', 298.15, 101325, 1e-6 * 6.02214076e+23),
+    (1, 'ppth', 298.15, 101325, 1e3 / calculate_air_density(298.15, 101325)),
+    (1, 'ppm', 298.15, 101325, 1e6 / calculate_air_density(298.15, 101325)),
+    (1, 'ppb', 298.15, 101325, 1e9 / calculate_air_density(298.15, 101325)),
+    (1, 'ppt', 298.15, 101325, 1e12 / calculate_air_density(298.15, 101325)),
+    (1, 'mol mol-1', 298.15, 101325, 1 / calculate_air_density(298.15, 101325)),
+])
+def test_convert_from_number_density(data, output_unit, temperature, pressure, expected):
+    assert math.isclose(convert_from_number_density(data, output_unit, temperature, pressure), expected)
