diff --git a/snsim/astrobj.py b/snsim/astrobj.py
index 8a18c9d..cb2ebb1 100644
--- a/snsim/astrobj.py
+++ b/snsim/astrobj.py
@@ -28,7 +28,7 @@ class AstrObj(abc.ABC):
     _obj_attrs = [""]
     _available_models = [""]
 
-    def __init__(self, sim_par, relation=None, effects=None):
+    def __init__(self, sim_par, mag_fun=None, effects=None):
         """Init AstrObj class.
 
         Parameters
@@ -45,8 +45,8 @@ def __init__(self, sim_par, relation=None, effects=None):
             | ├── dec, obj Declinaison
             | ├── t0, obj peak time
             | └── sncosmo par
-        relation : str, optional
-            _description_, by default None
+        mag_fun : str, optional
+            The function used to compute the abs mag, by default None
         effects : list, optional
             sncosmo effects dic, by default None
 
@@ -63,7 +63,7 @@ def __init__(self, sim_par, relation=None, effects=None):
         # -- Copy input parameters dic
         self._sim_par = copy.copy(sim_par)
 
-        self._relation = relation
+        self._mag_fun = mag_fun
 
         if "ID" not in self._sim_par:
             self._sim_par["ID"] = 0
@@ -194,31 +194,23 @@ def gen_flux(self, obs, mod_fcov=False, seed=None):
                 obs["band"], obs["time"], zp=obs["zp"], zpsys=obs["zpsys"]
             )
 
-            #compute the noise from the host galaxy if required
-            if "host_galaxy_noise" in self._sim_par:
-                if self._sim_par["host_galaxy_noise"]:
-                    sig_host = ut.model_galaxy_noise(self._sim_par,obs)
-                else:
-                    sig_host = np.zeros(len(fluxtrue))
-            else:
-                sig_host = np.zeros(len(fluxtrue))
+        sig_host = 0
+        #compute the Noise from the host galaxy if required
+        if "host_galaxy_noise" in self._sim_par:
+            if self._sim_par["host_galaxy_noise"]:
+                sig_host = ut.model_galaxy_noise(self._sim_par, obs)
 
         # -- Noise computation : Poisson Noise + Skynoise + ZP noise + Host gal Noise
         fluxerrtrue = np.sqrt(
-            np.abs(fluxtrue) / obs.gain
-            + obs.skynoise**2
-            + (np.log(10) / 2.5 * fluxtrue * obs.sig_zp) ** 2 
+            np.abs(fluxtrue) / obs['gain']
+            + obs['skynoise']**2
+            + (np.log(10) / 2.5 * fluxtrue * obs['sig_zp']) ** 2 
             + sig_host**2
         )
 
         gen = np.random.default_rng(random_seeds[1])
         flux = fluxtrue + gen.normal(loc=0.0, scale=fluxerrtrue)
-        fluxerr = np.sqrt(
-            np.abs(flux) / obs.gain
-            + obs.skynoise**2
-            + (np.log(10) / 2.5 * flux * obs.sig_zp) ** 2 
-            + sig_host**2
-        )
+        fluxerr = np.sqrt(fluxerrtrue**2 + (np.abs(flux) - np.abs(fluxtrue)) / obs['gain'])
 
         # -- Set magnitude
         mag = np.zeros_like(flux)
@@ -230,7 +222,7 @@ def gen_flux(self, obs, mod_fcov=False, seed=None):
         mag[positive_fmask] = -2.5 * np.log10(flux_pos) + obs["zp"][positive_fmask]
 
         magerr[positive_fmask] = (
-            2.5 / np.log(10) * 1 / flux_pos * fluxerr[positive_fmask]
+            2.5 / np.log(10) * fluxerr[positive_fmask] / flux_pos
         )
 
         mag[~positive_fmask] = np.nan
@@ -348,17 +340,17 @@ def _set_model_par(self, model):
         Raises
         ------
         ValueError
-            Raises if you use relation 'salttripp' for a non salt model.
+            Raises if you use mag_fun 'salttripp' for a non salt model.
         ValueError
-            Raises if you use a non-implemented relation.
+            Raises if you use a non-implemented mag_fun.
         ValueError
             Raises if you use mass-step without host logmass.
         """
 
-        if self._relation is None:
-            self._relation = "salttripp"
+        if self._mag_fun is None:
+            self._mag_fun = "salttripp"
 
-        if self._relation.lower() == "salttripp":
+        if self._mag_fun.lower() == "salttripp":
             if model.source.name not in ["salt2", "salt3"]:
                 raise ValueError("SALTTripp only available for salt2 & salt3 models")
 
@@ -378,8 +370,8 @@ def _set_model_par(self, model):
                 + self.mu
             )
         else:
-            # TODO - BC : Find a way to use lambda function for relation
-            raise ValueError("Relation not available")
+            # TODO - BC : Find a way to use lambda function for mag_fun
+            raise ValueError("mag_fun not available")
 
         if "mass_step" in self._sim_par:
             if "host_mass" in self._sim_par:
diff --git a/snsim/generators.py b/snsim/generators.py
index f52e110..cca7826 100644
--- a/snsim/generators.py
+++ b/snsim/generators.py
@@ -142,11 +142,6 @@ def __init__(
         # -- Get the astrobj class
         self._astrobj_class = getattr(astr, self._object_type)
 
-        # if peak_out_trange:
-        #     t0 = self.time_range[0] - self.snc_model_time[1] * (1 + self.z_range[1])
-        #     t1 = self.time_range[1] - self.snc_model_time[0] * (1 + self.z_range[1])
-        #     self._time_range = (t0, t1)
-
     def __call__(self, n_obj=None, seed=None, basic_par=None):
         """Launch the simulation of obj.
 
@@ -182,25 +177,23 @@ def __call__(self, n_obj=None, seed=None, basic_par=None):
         random_models = self.random_models(n_obj, seed=seeds[2])
 
         # -- Check if there is dust
+        dust_par = {}
         if self.mw_dust is not None:
             dust_par = self._compute_dust_par(basic_par["ra"], basic_par["dec"])
-        else:
-            dust_par = {}
 
         par = pd.DataFrame(
             {**random_models, **obj_par, **dust_par}, index=basic_par.index
-        )
+            )
 
         par = pd.concat([basic_par, par], axis=1)
 
-        if "relation" not in self._params:
-            relation = None
-        else:
-            relation = self._params["relation"]
+        mag_fun = None
+        if "mag_fun" in self._params :
+            mag_fun = self._params["mag_fun"]
 
         # TODO - BC: Dask that part or vectorize it for more efficiency
         return [
-            self._astrobj_class(par_dic, effects=self.sim_effects, relation=relation)
+            self._astrobj_class(par_dic, effects=self.sim_effects, mag_fun=mag_fun)
             for par_dic in par.reset_index().to_dict(orient="records")
         ]
 
diff --git a/snsim/scatter.py b/snsim/scatter.py
index fdd3172..e3a19c0 100644
--- a/snsim/scatter.py
+++ b/snsim/scatter.py
@@ -144,9 +144,9 @@ def propagate(self, wave, flux):
         return flux * 10 ** (-0.4 * magscat)
 
 
-##########################################
+###########################################
 # GENERATE terms for BS20 scattering model#
-##########################################
+###########################################
 def gen_BS20_scatter(n_sn, par_names=['beta_sn', 'Rv', 'E_dust', 'c_int'], seed=None):
     """Generate n coherent mag scattering term.
 
diff --git a/snsim/survey_host.py b/snsim/survey_host.py
index 6416e80..23a9a03 100644
--- a/snsim/survey_host.py
+++ b/snsim/survey_host.py
@@ -277,7 +277,7 @@ def _init_data(self):
             obs_dic["sig_zp"] = self.zp[1]
 
         if self.fwhm_psf != "psf_in_obs":
-            obs_dic["fwhm_psf"] = self.fwhm_psf
+            obs_dic["fwhm_psf"] = self.fwhm_psf 
 
         if self.gain != "gain_in_obs":
             obs_dic["gain"] = self.gain
@@ -539,12 +539,14 @@ def _make_obs_table(self, Obs):
             Obs["skynoise"] = Obs[self.config["noise_key"][0]]
         else:
             raise ValueError("Noise type should be mlim5 or skysigADU")
-    
+
+        # Add CCD noise
         if "ccd_noise" in self.config:
             Obs["skynoise"] = np.sqrt(Obs["skynoise"]**2 + self.config["ccd_noise"]**2)
-            
+
         # Apply PSF
-        Obs["skynoise"] *= np.sqrt(4 * np.pi) * Obs["fwhm_psf"] / (2 * np.sqrt(2 * np.log(2)))
+        if self.fwhm_psf != 0:
+            Obs["skynoise"] *= np.sqrt(4 * np.pi) * Obs["fwhm_psf"] / (2 * np.sqrt(2 * np.log(2)))
 
         # Magnitude system
         Obs["zpsys"] = "ab"
@@ -773,6 +775,7 @@ def compute_weights(self, rate=None, sn_type=None, cosmology=None):
         elif rate is not None:
             # Take into account rate is divide by (1 + z)
             weights_rate = rate(self.table["zcos"]) / (1 + self.table["zcos"])
+            
             if self.config["distrib"].lower() == "rate":
                 weights = weights_rate
             elif self.config["distrib"].lower() == "mass":
diff --git a/snsim/utils.py b/snsim/utils.py
index c82aa76..b9f4f39 100644
--- a/snsim/utils.py
+++ b/snsim/utils.py
@@ -704,25 +704,27 @@ def compute_weight_mass_sfr_for_type(mass,sfr, sn_type, cosmology):
     return weights
 
 
+def model_galaxy_noise(sim_par, obs):
+    """Compute noise coming from host galaxy in photoelectron units (approximation)"""
 
-def model_galaxy_noise(sim_par,obs):
-    """ compute noise coming from host galaxy in photoelectron units (approximation)"""
-
-    bands = ["host_"+ b for b in obs['band']]
-    mag_host = np.asarray([sim_par[bb] for bb in bands])
+    bands = ["host_" + b for b in obs['band']]
+    mag_host = np.asarray([sim_par[b] for b in bands])
 
     #compute mean surface brightness of the galaxy
     if sim_par.keys() & {"host_a_sersic", 'host_b_sersic'}:
-        if not isinstance(sim_par["host_a_sersic"], float ):
+        if not isinstance(sim_par["host_a_sersic"], float):
             if "host_w_sersic" not in sim_par.keys():
                 raise ValueError('if you have 2 or more sersic profile for each galaxy, Provide the weights')
             else:
-                s = np.asarray([mag_host  / (w* a * b ) for w,a,b 
-                in zip(sim_par['host_w_sersic'],sim_par['host_a_sersic'],sim_par['host_b_sersic'])])
+                s = np.asarray([mag_host  / (w * a * b) for w, a, b 
+                in zip(sim_par['host_w_sersic'],
+                       sim_par['host_a_sersic'],
+                       sim_par['host_b_sersic'])])
+                
                 surf_bright = np.sum(s,axis=0)
 
         elif isinstance(sim_par["host_a_sersic"], float):
-            surf_bright = mag_host  / (  sim_par['host_a_sersic'] * sim_par['host_b_sersic'] )
+            surf_bright = mag_host  / (sim_par['host_a_sersic'] * sim_par['host_b_sersic'])
     
         else:
             raise ValueError('for sersic parameters provide list or float for each host galaxies')
@@ -731,7 +733,7 @@ def model_galaxy_noise(sim_par,obs):
         raise ValueError('provide sersic ellipse parameters as host_a_sersic & host_b_sersic')
         
     # compute galaxy flux at SN position
-    mm = surf_bright * np.pi * 4 * obs.sig_psf **2
-    F = 10.**(0.4 * (obs.zp - mm))
+    mm = surf_bright * np.pi * 4 * obs['sig_psf'] **2
+    F = 10.**(0.4 * (obs['zp'] - mm))
 
-    return np.abs(F) / obs.gain
\ No newline at end of file
+    return np.abs(F) / obs['gain']
\ No newline at end of file