From 41e5ecabf780f7a28355f0dcd1707ff6e5aafe47 Mon Sep 17 00:00:00 2001
From: Satwik Kambham <satwik.kambham@gmail.com>
Date: Mon, 8 Aug 2022 20:08:39 +0530
Subject: [PATCH] Added atomic mass unit as a keyword argument

---
 tardis/model/base.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tardis/model/base.py b/tardis/model/base.py
index 494fc803378..f4fce247b0a 100644
--- a/tardis/model/base.py
+++ b/tardis/model/base.py
@@ -94,6 +94,7 @@ 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
     ----------
@@ -101,10 +102,16 @@ class Composition:
         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