Upgrade Earlinet reader and implement vertical profiles

pyaerocom/aeroval/coldatatojson_engine.py

@@ -1,7 +1,10 @@
 import logging
 import os
+import shutil
 from time import time
 
+from cf_units import Unit
+
 from pyaerocom import ColocatedData, TsType
 from pyaerocom._lowlevel_helpers import write_json
 from pyaerocom.aeroval._processing_base import ProcessingEngine
@@ -18,10 +21,13 @@
     _process_statistics_timeseries,
     _write_site_data,
     _write_stationdata_json,
+    add_profile_entry_json,
     get_heatmap_filename,
     get_json_mapname,
+    get_profile_filename,
     get_timeseries_file_name,
     init_regions_web,
+    process_profile_data,
     update_regions_json,
 )
 from pyaerocom.exceptions import AeroValConfigError, TemporalResolutionError
@@ -46,6 +52,11 @@
             list of files that have been converted.
 
         """
+        out_dirs = self.cfg.path_manager.get_json_output_dirs(True)
+        for idir in out_dirs:
+            if os.path.exists(out_dirs[idir]):
+                shutil.rmtree(out_dirs[idir])
+
         converted = []
         for file in files:
             logger.info(f"Processing: {file}")
@@ -93,7 +104,34 @@
 
         stats_min_num = self.cfg.statistics_opts.MIN_NUM
 
-        vert_code = coldata.get_meta_item("vert_code")
+        if hasattr(coldata.data, "vertical_layer"):
+            if Unit(coldata.data.attrs["altitude_units"]) != Unit(
+                "km"
+            ):  # put everything in terms of km for viz
+                # convert start and end for file naming
+                alt_units = coldata.data.attrs["altitude_units"]
+
+                coldata.data.attrs["vertical_layer"]["start"] = str(
+                    Unit(alt_units).convert(
+                        coldata.data.attrs["vertical_layer"]["start"], other="km"
+                    )
+                )
+
+                coldata.data.attrs["vertical_layer"]["end"] = str(
+                    Unit(alt_units).convert(
+                        coldata.data.attrs["vertical_layer"]["end"], other="km"
+                    )
+                )
+                # start and end for vertical layers (display on web and name jsons)
+                start = float(coldata.data.attrs["vertical_layer"]["start"])
+                end = float(coldata.data.attrs["vertical_layer"]["end"])
+                # format correctly (e.g., 1, 1.5, 2, 2.5, etc.)
+                start = f"{round(float(start), 1):g}"
+                end = f"{round(float(end), 1):g}"
+                vert_code = f"{start}-{end}km"
+        else:
+            vert_code = coldata.get_meta_item("vert_code")
+
         diurnal_only = coldata.get_meta_item("diurnal_only")
 
         add_trends = self.cfg.statistics_opts.add_trends
@@ -109,7 +147,7 @@
 
         # this will need to be figured out as soon as there is altitude
         elif "altitude" in coldata.data.dims:
-            raise NotImplementedError("Cannot yet handle profile data")
+            raise ValueError("Altitude should have been dealt with already in the colocation")
 
         elif not isinstance(coldata, ColocatedData):
             raise ValueError(f"Need ColocatedData object, got {type(coldata)}")
@@ -165,114 +203,159 @@
         if annual_stats_constrained:
             data = _apply_annual_constraint(data)
 
-        if not diurnal_only:
-            logger.info("Processing statistics timeseries for all regions")
-            input_freq = self.cfg.statistics_opts.stats_tseries_base_freq
-
-            for reg in regnames:
-                try:
-                    stats_ts = _process_statistics_timeseries(
-                        data=data,
-                        freq=main_freq,
-                        region_ids={reg: regnames[reg]},
-                        use_weights=use_weights,
-                        use_country=use_country,
-                        data_freq=input_freq,
+        if (
+            not "just_for_viz" in coldata.data.attrs or coldata.data.attrs["just_for_viz"] == 0
+        ):  # make the regular json output
+            if not diurnal_only:
+                logger.info("Processing statistics timeseries for all regions")
+                input_freq = self.cfg.statistics_opts.stats_tseries_base_freq
+                for reg in regnames:
+                    try:
+                        stats_ts = _process_statistics_timeseries(
+                            data=data,
+                            freq=main_freq,
+                            region_ids={reg: regnames[reg]},
+                            use_weights=use_weights,
+                            use_country=use_country,
+                            data_freq=input_freq,
+                        )
+
+                    except TemporalResolutionError:
+                        stats_ts = {}
+                    fname = get_timeseries_file_name(
+                        regnames[reg], obs_name, var_name_web, vert_code
+                    )
+                    ts_file = os.path.join(out_dirs["hm/ts"], fname)
+                    _add_heatmap_entry_json(
+                        ts_file, stats_ts, obs_name, var_name_web, vert_code, model_name, model_var
                     )
 
-                except TemporalResolutionError:
-                    stats_ts = {}
+                logger.info("Processing heatmap data for all regions")
 
-                fname = get_timeseries_file_name(regnames[reg], obs_name, var_name_web, vert_code)
-                ts_file = os.path.join(out_dirs["hm/ts"], fname)
-                _add_heatmap_entry_json(
-                    ts_file, stats_ts, obs_name, var_name_web, vert_code, model_name, model_var
+                hm_all = _process_heatmap_data(
+                    data,
+                    regnames,
+                    use_weights,
+                    use_country,
+                    meta_glob,
+                    periods,
+                    seasons,
+                    add_trends,
+                    trends_min_yrs,
                 )
 
-            logger.info("Processing heatmap data for all regions")
-            hm_all = _process_heatmap_data(
-                data,
-                regnames,
-                use_weights,
-                use_country,
-                meta_glob,
-                periods,
-                seasons,
-                add_trends,
-                trends_min_yrs,
-            )
+                for freq, hm_data in hm_all.items():
+                    fname = get_heatmap_filename(freq)
 
-            for freq, hm_data in hm_all.items():
-                fname = get_heatmap_filename(freq)
+                    hm_file = os.path.join(out_dirs["hm"], fname)
 
-                hm_file = os.path.join(out_dirs["hm"], fname)
+                    _add_heatmap_entry_json(
+                        hm_file, hm_data, obs_name, var_name_web, vert_code, model_name, model_var
+                    )
 
-                _add_heatmap_entry_json(
-                    hm_file, hm_data, obs_name, var_name_web, vert_code, model_name, model_var
+                logger.info("Processing regional timeseries for all regions")
+                ts_objs_regional = _process_regional_timeseries(
+                    data, regnames, regions_how, meta_glob
                 )
 
-            logger.info("Processing regional timeseries for all regions")
-            ts_objs_regional = _process_regional_timeseries(data, regnames, regions_how, meta_glob)
+                _write_site_data(ts_objs_regional, out_dirs["ts"])
+                if coldata.has_latlon_dims:
+                    for cd in data.values():
+                        if cd is not None:
+                            cd.data = cd.flatten_latlondim_station_name().data
 
-            _write_site_data(ts_objs_regional, out_dirs["ts"])
-            if coldata.has_latlon_dims:
-                for cd in data.values():
-                    if cd is not None:
-                        cd.data = cd.flatten_latlondim_station_name().data
+                logger.info("Processing individual site timeseries data")
+                (ts_objs, map_meta, site_indices) = _process_sites(
+                    data, regs, regions_how, meta_glob
+                )
 
-            logger.info("Processing individual site timeseries data")
-            (ts_objs, map_meta, site_indices) = _process_sites(data, regs, regions_how, meta_glob)
+                _write_site_data(ts_objs, out_dirs["ts"])
+
+                scatter_freq = min(TsType(fq) for fq in self.cfg.time_cfg.freqs)
+                scatter_freq = min(scatter_freq, main_freq)
+
+                logger.info("Processing map and scat data by period")
+
+                for period in periods:
+                    # compute map_data and scat_data just for this period
+                    map_data, scat_data = _process_map_and_scat(
+                        data,
+                        map_meta,
+                        site_indices,
+                        [period],
+                        str(scatter_freq),
+                        stats_min_num,
+                        seasons,
+                        add_trends,
+                        trends_min_yrs,
+                        use_fairmode,
+                        obs_var,
+                    )
 
-            _write_site_data(ts_objs, out_dirs["ts"])
+                    # the files in /map and /scat will be split up according to their time period as well
+                    map_name = get_json_mapname(
+                        obs_name, var_name_web, model_name, model_var, vert_code, period
+                    )
+                    outfile_map = os.path.join(out_dirs["map"], map_name)
+                    write_json(map_data, outfile_map, ignore_nan=True)
 
-            scatter_freq = min(TsType(fq) for fq in self.cfg.time_cfg.freqs)
-            scatter_freq = min(scatter_freq, main_freq)
+                    outfile_scat = os.path.join(out_dirs["scat"], map_name)
+                    write_json(scat_data, outfile_scat, ignore_nan=True)
 
-            logger.info("Processing map and scat data by period")
-            for period in periods:
-                # compute map_data and scat_data just for this period
-                map_data, scat_data = _process_map_and_scat(
-                    data,
-                    map_meta,
-                    site_indices,
-                    [period],
-                    str(scatter_freq),
-                    stats_min_num,
-                    seasons,
-                    add_trends,
-                    trends_min_yrs,
-                    use_fairmode,
-                    obs_var,
+            if coldata.ts_type == "hourly" and use_diurnal:
+                logger.info("Processing diurnal profiles")
+                (ts_objs_weekly, ts_objs_weekly_reg) = _process_sites_weekly_ts(
+                    coldata, regions_how, regnames, meta_glob
                 )
+                outdir = os.path.join(out_dirs["ts/diurnal"])
+                for ts_data_weekly in ts_objs_weekly:
+                    # writes json file
+                    _write_stationdata_json(ts_data_weekly, outdir)
+                if ts_objs_weekly_reg != None:
+                    for ts_data_weekly_reg in ts_objs_weekly_reg:
+                        # writes json file
+                        _write_stationdata_json(ts_data_weekly_reg, outdir)
+        else:
+            if hasattr(coldata.data, "vertical_layer"):
+                logger.info("Processing profile data for vizualization")
+                # Loop through regions
+                for regid in regnames:
+                    profile_viz = process_profile_data(
+                        data=data,
+                        region_id=regid,
+                        station_name=None,
+                        use_country=use_country,
+                        periods=periods,
+                        seasons=seasons,
+                    )
 
-                # the files in /map and /scat will be split up according to their time period as well
-                map_name = get_json_mapname(
-                    obs_name, var_name_web, model_name, model_var, vert_code, period
-                )
-                outfile_map = os.path.join(out_dirs["map"], map_name)
-                write_json(map_data, outfile_map, ignore_nan=True)
+                    fname = get_profile_filename(regnames[regid], obs_name, var_name_web)
+
+                    outfile_profile = os.path.join(out_dirs["profiles"], fname)
+                    add_profile_entry_json(outfile_profile, data, profile_viz, periods, seasons)
+                # Loop through stations
+                for station_name in coldata.data.station_name.values:
+                    profile_viz = process_profile_data(
+                        data=data,
+                        region_id=None,
+                        station_name=station_name,
+                        use_country=use_country,
+                        periods=periods,
+                        seasons=seasons,
+                    )
 
-                outfile_scat = os.path.join(out_dirs["scat"], map_name)
-                write_json(scat_data, outfile_scat, ignore_nan=True)
+                    fname = get_profile_filename(station_name, obs_name, var_name_web)
 
-        if coldata.ts_type == "hourly" and use_diurnal:
-            logger.info("Processing diurnal profiles")
-            (ts_objs_weekly, ts_objs_weekly_reg) = _process_sites_weekly_ts(
-                coldata, regions_how, regnames, meta_glob
-            )
-            outdir = os.path.join(out_dirs["ts/diurnal"])
-            for ts_data_weekly in ts_objs_weekly:
-                # writes json file
-                _write_stationdata_json(ts_data_weekly, outdir)
-            if ts_objs_weekly_reg != None:
-                for ts_data_weekly_reg in ts_objs_weekly_reg:
-                    # writes json file
-                    _write_stationdata_json(ts_data_weekly_reg, outdir)
+                    outfile_profile = os.path.join(out_dirs["profiles"], fname)
+                    add_profile_entry_json(outfile_profile, data, profile_viz, periods, seasons)
 
         logger.info(
             f"Finished computing json files for {model_name} ({model_var}) vs. "
             f"{obs_name} ({obs_var})"
         )
 
         dt = time() - t00
-        logger.info(f"Time expired (TOTAL): {dt:.2f} s")
+        logger.info(f"Time expired: {dt:.2f} s")
+
+        dt = time() - t00
+        logger.info(f"Time expired: {dt:.2f} s")