diff --git a/galaxy.py b/galaxy.py
index 3098f23..037004b 100644
--- a/galaxy.py
+++ b/galaxy.py
@@ -82,7 +82,7 @@ def extract_polar_region(dset, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
     mask_box = (
         (np.abs(dset["r"] - r) < dr)
         & (np.abs(dset["phi"] - phi) < dphi)
-        & (np.abs(dset["pos"][:, 2] - z) < dz)
+        & (np.abs(dset["pos_kpc"][:, 2] - z) < dz)
     )
     return {key: dset[key][mask_box] for key in dset}
 
@@ -143,14 +143,18 @@ def get_time_from_relax(pp):
     return tfromrelax
 
 
-def get_last_sfr(pp, r=4, dr=0.5):
+def get_last_sfr(pp, r=4, dr=0.5, z=0, dz=0.5):
     pp.load_parts()
     try:
-        epoch = pp.parts["epoch"].copy()
+        epoch = pp.stars["epoch"].copy()
         epoch *= pp.info["unit_time"].express(U.year)
-        mass = pp.parts["mass"].copy()
+        mass = pp.stars["mass"].copy()
         mass *= pp.info["unit_mass"].express(U.Msun)
-        mask = (epoch > 0) & (np.abs(pp.parts["r"] - r) < dr)
+        mask = (
+            (epoch > 0)
+            & (np.abs(pp.stars["r"] - r) < dr)
+            & (np.abs(pp.stars["pos_kpc"][:, 2] - z) < dz)
+        )
         masstot, time = np.histogram(epoch[mask], weights=mass[mask], bins=100)
         dtime = np.diff(time)
         sfr = masstot[-1] / dtime[-1]