fix output models

romancal/ramp_fitting/ramp_fit_step.py

@@ -19,46 +19,13 @@
 __all__ = ["RampFitStep"]
 
 
-def create_integration_model(input_model, integ_info):
-    """
-    Creates an ImageModel from the computed arrays from ramp_fit.
-    Parameter
-    ---------
-    input_model: RampModel
-        Input RampModel for which the output CubeModel is created.
-    integ_info: tuple
-        The ramp fitting arrays needed for the CubeModel for each integration.
-    Parameter
-    ---------
-    int_model: CubeModel
-        The output CubeModel to be returned from the ramp fit step.
-    """
-    data, dq, var_poisson, var_rnoise, int_times, err = integ_info
-    int_model = rdd.CubeModel(
-        data=np.zeros(data.shape, dtype=np.float32),
-        dq=np.zeros(data.shape, dtype=np.uint32),
-        var_poisson=np.zeros(data.shape, dtype=np.float32),
-        var_rnoise=np.zeros(data.shape, dtype=np.float32),
-        err=np.zeros(data.shape, dtype=np.float32))
-    int_model.int_times = None
-    int_model.update(input_model)  # ... and add all keys from input
-
-    int_model.data = data
-    int_model.dq = dq
-    int_model.var_poisson = var_poisson
-    int_model.var_rnoise = var_rnoise
-    int_model.err = err
-    int_model.int_times = int_times
-
-    return int_model
-
-
 def create_optional_results_model(input_model, opt_info):
     """
     Creates an ImageModel from the computed arrays from ramp_fit.
-    Parameter
-    ---------
-    input_model: ~jwst.datamodels.RampModel
+
+    Parameters
+    ----------
+    input_model: ~roman_datamodels.RampModel
     opt_info: tuple
         The ramp fitting arrays needed for the RampFitOutputModel.
     Parameter
@@ -68,36 +35,44 @@ def create_optional_results_model(input_model, opt_info):
     """
     (slope, sigslope, var_poisson, var_rnoise,
         yint, sigyint, pedestal, weights, crmag) = opt_info
-    opt_model = rdd.RampFitOutputModel(
-        slope=slope,
-        sigslope=sigslope,
-        var_poisson=var_poisson,
-        var_rnoise=var_rnoise,
-        yint=yint,
-        sigyint=sigyint,
-        pedestal=pedestal,
-        weights=weights,
-        crmag=crmag)
+    meta = {}
+    meta.update(input_model.meta)
+    crmag.shape = crmag.shape[1:]
+    crmag.dtype = np.float32
 
-    opt_model.meta.filename = input_model.meta.filename
-    opt_model.update(input_model)  # ... and add all keys from input
+    inst = {'meta': meta,
+            'slope': np.squeeze(slope),
+            'sigslope': np.squeeze(sigslope),
+            'var_poisson': np.squeeze(var_poisson),
+            'var_rnoise': np.squeeze(var_rnoise),
+            'yint': np.squeeze(yint),
+            'sigyint': np.squeeze(sigyint),
+            'pedestal': np.squeeze(pedestal),
+            'weights': np.squeeze(weights),
+            'crmag': crmag
+            }
 
+    out_node = rds.RampFitOutput(inst)
+    opt_model = rdd.RampFitOutputModel(out_node)
+    opt_model.meta.filename = input_model.meta.filename
     return opt_model
 
 
 def create_image_model(input_model, image_info):
     """
     Creates an ImageModel from the computed arrays from ramp_fit.
-    Parameter
-    ---------
-    input_model: RampModel
-        Input RampModel for which the output ImageModel is created.
+
+    Parameters
+    ----------
+    input_model : `~roman_datamodels.RampModel`
+        Input ``RampModel`` for which the output ImageModel is created.
     image_info: tuple
         The ramp fitting arrays needed for the ImageModel.
-    Parameter
-    ---------
-    out_model: ImageModel
-        The output ImageModel to be returned from the ramp fit step.
+
+    Returns
+    -------
+    out_model : `~roman_datamodels.ImageModel`
+        The output ``ImageModel`` to be returned from the ramp fit step.
     """
     data, dq, var_poisson, var_rnoise, err = image_info
 
@@ -111,14 +86,11 @@ def create_image_model(input_model, image_info):
             'dq': dq,
             'var_poisson': var_poisson,
             'var_rnoise': var_rnoise,
-            'err': err
+            'err': err,
             }
-    out_model = rds.WfiImage({'roman': inst})
-    af = asdf.AsdfFile()
-    af.tree = {'roman': out_model}
-    # return out_model
-    return af
-
+    out_node = rds.WfiImage(inst)
+    im = rdd.ImageModel(out_node)
+    return im
 
 
 class RampFitStep(RomanStep):
@@ -133,10 +105,9 @@ class RampFitStep(RomanStep):
         maximum_cores = option('none','quarter','half','all',default='none') # max number of processes to create
         save_opt = boolean(default=False) # Save optional output
     """
-    algorithm = 'ols'      # Only algorithm allowed ?
-#    algorithm = 'gls'
+    algorithm = 'ols'      # Only algorithm allowed
 
-    weighting = 'optimal'  # Only weighting allowed ?
+    weighting = 'optimal'  # Only weighting allowed
 
     reference_file_types = ['readnoise', 'gain']
 
@@ -159,20 +130,10 @@ def process(self, input):
             log.info('Using weighting = %s' % self.weighting)
 
             buffsize = ramp_fit.BUFSIZE
-
             image_info, integ_info, opt_info, gls_opt_model = ramp_fit.ramp_fit(
                 input_model, buffsize, self.save_opt,
                 readnoise_model.data, gain_model.data, self.algorithm,
                 self.weighting, max_cores, dqflags.pixel)
-            # The above int_model is None
-
-            # Remove the empty 0th dimension from the JWST model and insert
-            # the results into the Roman DataModel
-            #input_model.data = input_model.data[-1]
-            #R_input_model.groupdq = input_model.groupdq[-1]
-            #R_input_model.err = input_model.err[-1]
-            #R_input_model.refout = input_model.refout[-1]
-
             readnoise_model.close()
             gain_model.close()
 
@@ -183,11 +144,6 @@ def process(self, input):
 
         if image_info is not None:
             out_model = create_image_model(input_model, image_info)
-            #out_model.meta.cal_step.ramp_fit = 'COMPLETE'
-
-        # if integ_info is not None:
-        #     log.info(f'integration info, {integ_info}')
-        #     int_model = create_integration_model(input_model, integ_info)
-        #     int_model.meta.cal_step.ramp_fit = 'COMPLETE'
+            out_model.meta.cal_step.ramp_fit = 'COMPLETE'
 
         return out_model, None  # 'None' for int_model