Added atomic mass unit as a keyword argument

tardis/model/base.py

@@ -94,17 +94,24 @@ class Composition:
         An array of densities for each shell.
     isotopic_mass_fraction : pd.DataFrame
     atomic_mass : pandas.core.series.Series
+    atomic_mass_unit: astropy.units.Unit
 
     Attributes
     ----------
     number_density : pd.DataFrame
         Number density of each isotope in each shell.
     """
 
-    def __init__(self, density, elemental_mass_fraction, atomic_mass):
+    def __init__(
+        self,
+        density,
+        elemental_mass_fraction,
+        atomic_mass,
+        atomic_mass_unit=u.g,
+    ):
         self.density = density
         self.elemental_mass_fraction = elemental_mass_fraction
-        self.atomic_mass_unit = u.g
+        self.atomic_mass_unit = atomic_mass_unit
         self.atomic_mass = atomic_mass
 
     @property