def calc()

in analysis/webservice/algorithms/Tomogram.py [0:0]

• 92 lines of code
• 22 McCabe index (conditional complexity)

    def calc(self, compute_options, **args):
        (dataset, parameter, latitude, min_lon, max_lon, min_elevation, max_elevation, horizontal_margin, calc_peaks,
         ch_ds, g_ds, percentiles, cmap) = self.parse_args(compute_options)

        slices = dict(
            lon=slice(min_lon, max_lon),
            elevation=slice(min_elevation, max_elevation)
        )

        if isinstance(latitude, slice):
            slices['lat'] = latitude
        elif isinstance(latitude, list):
            slices['lat'] = slice(latitude[0] - horizontal_margin, latitude[-1] + horizontal_margin)
        else:
            slices['lat'] = slice(latitude - horizontal_margin, latitude + horizontal_margin)

        data_in_bounds = self.do_subset(dataset, parameter, slices, 0)
        z_step = TomogramBaseClass._get_elevation_step_size(data_in_bounds)

        r = np.arange(min_elevation, max_elevation + z_step, z_step)

        logger.info(f'Fetched {len(data_in_bounds):,} data points at this elevation')

        ds = TomogramBaseClass.data_subset_to_ds_with_elevation(
            TomogramBaseClass.bin_subset_elevations_to_range(data_in_bounds, r, z_step / 2)
        )[3]

        slices['elevation'] = slice(None, None)
        elev_vars = {}

        if len(ch_ds) > 0:
            elev_vars['ch'] = TomogramBaseClass.data_subset_to_ds(
                self.do_subset(ch_ds[0], parameter, slices, 0)
            )[3]
            elev_vars['ch'].attrs['_source'] = ch_ds[0]

            if len(ch_ds) == 2:
                elev_vars['ch_secondary'] = TomogramBaseClass.data_subset_to_ds(
                    self.do_subset(ch_ds[1], parameter, slices, 0)
                )[3]
                elev_vars['ch_secondary'].attrs['_source'] = ch_ds[1]

        if len(g_ds) > 0:
            elev_vars['gh'] = TomogramBaseClass.data_subset_to_ds(
                self.do_subset(g_ds[0], parameter, slices, 0)
            )[3]
            elev_vars['gh'].attrs['_source'] = g_ds[0]

            if len(g_ds) == 2:
                elev_vars['gh_secondary'] = TomogramBaseClass.data_subset_to_ds(
                    self.do_subset(g_ds[1], parameter, slices, 0)
                )[3]
                elev_vars['gh_secondary'].attrs['_source'] = g_ds[1]

        if elev_vars:
            TomogramBaseClass.add_variables_to_tomo(ds, **elev_vars)

        if isinstance(latitude, list):
            lat_slices = [ds.sel(latitude=l, method='nearest') for l in latitude]
            ds = xr.concat(lat_slices, dim='latitude').drop_duplicates('latitude')
        else:
            ds = ds.sel(latitude=latitude)

        if 'ch' in ds.data_vars:
            if 'ch_secondary' in ds.data_vars:
                ds['ch_filter'] = xr.where(np.isnan(ds.ch), ds.ch_secondary, ds.ch)
            else:
                ds['ch_filter'] = ds.ch

            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation <= ds.ch_filter)

        if 'gh' in ds.data_vars:
            if 'gh_secondary' in ds.data_vars:
                ds['gh_filter'] = xr.where(np.isnan(ds.gh), ds.gh_secondary, ds.gh)
            else:
                ds['gh_filter'] = ds.gh

            ds['tomo'] = ds['tomo'].where(ds.tomo.elevation >= ds.gh_filter)

        if percentiles:
            tomo = ds['tomo'].cumsum(dim='elevation', skipna=True).where(ds['tomo'].notnull())
            ds['tomo'] = tomo / tomo.max(dim='elevation', skipna=True)
        else:
            ds['tomo'] = xr.apply_ufunc(lambda a: 10 * np.log10(a), ds.tomo)

        if calc_peaks:
            peaks = []

            for i in range(len(ds.tomo.latitude)):
                lat_peaks = []
                ds_sub = ds.isel(latitude=i)

                for j in range(len(ds.tomo.longitude)):
                    col = ds_sub.isel(longitude=j)

                    lon = col.longitude.item()

                    try:
                        idx = col.tomo.argmax(dim='elevation').item()

                        lat_peaks.append((lon, col.isel(elevation=idx).elevation.item()))

                    except ValueError:
                        lat_peaks.append((lon, np.nan))

                peaks.append((np.array([p[0] for p in lat_peaks]),
                              np.array([p[1] for p in lat_peaks])))
        else:
            peaks = None

        return ProfileTomoResults(
            results=(ds, peaks),
            s={'latitude': latitude},
            coords=dict(x=ds.longitude, y=ds.elevation),
            meta=dict(dataset=dataset),
            style='db' if not percentiles else 'percentile',
            cmap=cmap,
            computeOptions=compute_options
        )