solved #434 and #582: df column units

lasio/las.py

@@ -901,25 +901,46 @@ def header(self):
         """
         return self.sections
 
-    def df(self):
+    def df(self, include_units: bool=False):
         """Return data as a :class:`pandas.DataFrame` structure.
 
         The first Curve of the LASFile object is used as the pandas
         DataFrame's index.
 
+        Arguments:
+            include_units (bool): Set to `True` to include the unit in the column name(s) of the dataframe. Defaults to `False`.
+
+        Returns:
+            pandas.DataFrame
+
         """
         import pandas as pd
         from pandas.api.types import is_object_dtype
 
-        df = pd.DataFrame(self.data, columns=[c.mnemonic for c in self.curves])
+        def col_name(curve: CurveItem, include_units: bool) -> str:
+            """Returns a formatted column name, including the curves' unit
+            if include_units set to True.
+
+            Arguments:
+                curve (CurveItem): curve column.
+                include_units (bool): Set to True to include the unit in the column name of the dataframe.
+
+            Returns:
+                str
+            """
+            unit = f" ({curve.unit})" if include_units and curve.unit else ""
+            mnemonic = curve.original_mnemonic if include_units else curve.mnemonic
+            return f"{mnemonic}{unit}"
+
+        df = pd.DataFrame(self.data, columns=[col_name(c, include_units) for c in self.curves])
         for column in df.columns:
             if is_object_dtype(df[column].dtype):
                 try:
                     df[column] = df[column].astype(np.float64)
                 except ValueError:
                     pass
         if len(self.curves) > 0:
-            df = df.set_index(self.curves[0].mnemonic)
+            df = df.set_index(df.columns[0])
         return df
 
     @property

tests/examples/sample_curve_units.las

@@ -0,0 +1,46 @@
+~VERSION INFORMATION
+ VERS.                  1.2:   CWLS LOG ASCII STANDARD -VERSION 1.2
+ WRAP.                  NO:   ONE LINE PER DEPTH STEP
+~WELL INFORMATION BLOCK
+#MNEM.UNIT       DATA TYPE    INFORMATION
+#---------    -------------   ------------------------------
+ STRT.M        1670.000000:
+ STOP.M        1669.750000:
+ STEP.M            -0.1250:
+ NULL.           -999.2500:
+ COMP.             COMPANY:   # ANY OIL COMPANY LTD.
+ WELL.                WELL:   ANY ET AL OIL WELL #12
+ FLD .               FIELD:   EDAM
+ LOC .            LOCATION:   A9-16-49-20W3M
+ PROV.            PROVINCE:   SASKATCHEWAN
+ SRVC.     SERVICE COMPANY:   ANY LOGGING COMPANY LTD.
+ DATE.            LOG DATE:   25-DEC-1988
+ UWI .      UNIQUE WELL ID:   100091604920W300
+~CURVE INFORMATION
+#MNEM.UNIT      API CODE      CURVE DESCRIPTION
+#---------    -------------   ------------------------------
+ DEPT.M                      :  1  DEPTH
+ DT  .US/M                   :  2  SONIC TRANSIT TIME
+ RHOB.K/M3                   :  3  BULK DENSITY
+ NPHI.V/V                    :  4   NEUTRON POROSITY
+ NO_UNIT.                    :  5  NO UNIT TEST CASE
+ SFLA.OHMM                   :  6  SHALLOW RESISTIVITY
+ ILM .OHMM1                  :  7  MEDIUM RESISTIVITY
+ ILM .OHMM2                  :  8  MEDIUM RESISTIVITY
+~PARAMETER INFORMATION
+#MNEM.UNIT        VALUE       DESCRIPTION
+#---------    -------------   ------------------------------
+ BHT .DEGC         35.5000:   BOTTOM HOLE TEMPERATURE
+ BS  .MM          200.0000:   BIT SIZE
+ FD  .K/M3       1000.0000:   FLUID DENSITY
+ MATR.              0.0000:   NEUTRON MATRIX(0=LIME,1=SAND,2=DOLO)
+ MDEN.           2710.0000:   LOGGING MATRIX DENSITY
+ RMF .OHMM          0.2160:   MUD FILTRATE RESISTIVITY
+ DFD .K/M3       1525.0000:   DRILL FLUID DENSITY
+~Other
+     Note: The logging tools became stuck at 625 meters causing the data
+	   between 625 meters and 615 meters to be invalid.
+~A  DEPTH     DT       RHOB     NPHI     SFLU     SFLA      ILM      ILD
+1670.000   123.450 2550.000    0.450  123.450  123.450  110.200  105.600
+1669.875   123.450 2550.000    0.450  123.450  123.450  110.200  105.600
+1669.750   123.450 2550.000    0.450  123.450  123.450  110.200  105.600

tests/test_read.py

@@ -542,3 +542,18 @@ def test_tab_dlm_numpy_engine():
     # GitHub Issue 554
     las = lasio.examples.open("2.0/sample_2.0_tab_dlm.las", engine="numpy")
     assert las["DEPT"].data[1] == 1669.875
+
+
+def test_df_include_units():
+    # GitHub Issues #434 and #582
+    las = lasio.examples.open("sample_curve_units.las")
+    #434
+    assert las.df().columns[0] == 'DT'
+    assert las.df(include_units=True).columns[0] == 'DT (US/M)'
+    assert las.df().columns[3] == 'NO_UNIT'
+    assert las.df(include_units=True).columns[3] == 'NO_UNIT'
+    #582
+    assert las.df().columns[5] == 'ILM:1'
+    assert las.df().columns[6] == 'ILM:2'
+    assert las.df(include_units=True).columns[5] == 'ILM (OHMM1)'
+    assert las.df(include_units=True).columns[6] == 'ILM (OHMM2)'