feat: improve cluster shutdown and cleanup logic
This commit is contained in:
+312
-294
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import yaml
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
import dask.dataframe as dd
|
||||
import pandas as pd
|
||||
@@ -39,171 +41,139 @@ def main():
|
||||
args = parse_args()
|
||||
run_config = load_and_resolve_config(args)
|
||||
|
||||
client = initialize_cluster(run_config)
|
||||
client, cluster = initialize_cluster(run_config)
|
||||
|
||||
# -1. chunking
|
||||
pbp_times = extract_partitioned_datetimes(run_config["input_pbp"])
|
||||
hk_times = extract_partitioned_datetimes(run_config["input_hk"])
|
||||
global_start = min(min(pbp_times), min(hk_times))
|
||||
global_end = max(max(pbp_times), max(hk_times))
|
||||
chunk_freq = run_config["chunking"]["freq"] # e.g. "6h", "3d"
|
||||
time_chunks = get_time_chunks_from_range(global_start, global_end, chunk_freq)
|
||||
|
||||
# 0. calibration stage --------------------------------------------
|
||||
instr_config = yaml.safe_load(open(run_config["instr_cfg"]))
|
||||
|
||||
# 1. Bins
|
||||
inc_mass_bin_lims = np.logspace(
|
||||
np.log10(run_config["histo"]["inc"]["min_mass"]),
|
||||
np.log10(run_config["histo"]["inc"]["max_mass"]),
|
||||
run_config["histo"]["inc"]["n_bins"],
|
||||
)
|
||||
inc_mass_bin_ctrs = bin_lims_to_ctrs(inc_mass_bin_lims)
|
||||
|
||||
scatt_bin_lims = np.logspace(
|
||||
np.log10(run_config["histo"]["scatt"]["min_D"]),
|
||||
np.log10(run_config["histo"]["scatt"]["max_D"]),
|
||||
run_config["histo"]["scatt"]["n_bins"],
|
||||
)
|
||||
scatt_bin_ctrs = bin_lims_to_ctrs(scatt_bin_lims)
|
||||
|
||||
timelag_bins_lims = np.linspace(
|
||||
run_config["histo"]["timelag"]["min"],
|
||||
run_config["histo"]["timelag"]["max"],
|
||||
run_config["histo"]["timelag"]["n_bins"],
|
||||
)
|
||||
timelag_bin_ctrs = bin_lims_to_ctrs(timelag_bins_lims)
|
||||
|
||||
for chunk_start, chunk_end in time_chunks:
|
||||
print(f"Processing: {chunk_start} to {chunk_end}")
|
||||
|
||||
pbp_filters = [
|
||||
("date", ">=", chunk_start.date().strftime("%Y-%m-%d")),
|
||||
("date", "<", chunk_end.date().strftime("%Y-%m-%d")),
|
||||
]
|
||||
if "hour" in run_config["chunking"]["freq"]: # optionally filter by hour
|
||||
pbp_filters.append(("hour", ">=", chunk_start.hour))
|
||||
pbp_filters.append(("hour", "<", chunk_end.hour))
|
||||
|
||||
# 2. HK processing --------------------------------------
|
||||
def handle_sigterm(signum, frame):
|
||||
print(
|
||||
f"\nSIGTERM received (signal {signum}), shutting down Dask...", flush=True
|
||||
)
|
||||
try:
|
||||
ddf_hk = dd.read_parquet(
|
||||
run_config["input_hk"],
|
||||
engine="pyarrow",
|
||||
filters=pbp_filters,
|
||||
calculate_divisions=True,
|
||||
)
|
||||
except (FileNotFoundError, OSError):
|
||||
print(" → no HK files for this chunk; skipping.")
|
||||
continue
|
||||
client.close()
|
||||
cluster.close()
|
||||
except Exception as e:
|
||||
print(f"Error during cleanup: {e}", flush=True)
|
||||
sys.exit(0)
|
||||
|
||||
if ddf_hk.npartitions == 0 or partition_rowcount(ddf_hk) == 0:
|
||||
print(" → HK frame is empty; skipping.")
|
||||
continue
|
||||
ddf_hk = ddf_hk.map_partitions(lambda pdf: pdf.sort_index())
|
||||
if not ddf_hk.known_divisions:
|
||||
ddf_hk = (
|
||||
ddf_hk.reset_index().set_index( # 'calculated_time' becomes a column
|
||||
signal.signal(signal.SIGTERM, handle_sigterm)
|
||||
try:
|
||||
# -1. chunking
|
||||
pbp_times = extract_partitioned_datetimes(run_config["input_pbp"])
|
||||
hk_times = extract_partitioned_datetimes(run_config["input_hk"])
|
||||
global_start = min(min(pbp_times), min(hk_times))
|
||||
global_end = max(max(pbp_times), max(hk_times))
|
||||
chunk_freq = run_config["chunking"]["freq"] # e.g. "6h", "3d"
|
||||
time_chunks = get_time_chunks_from_range(global_start, global_end, chunk_freq)
|
||||
|
||||
# 0. calibration stage --------------------------------------------
|
||||
instr_config = yaml.safe_load(open(run_config["instr_cfg"]))
|
||||
|
||||
# 1. Bins
|
||||
inc_mass_bin_lims = np.logspace(
|
||||
np.log10(run_config["histo"]["inc"]["min_mass"]),
|
||||
np.log10(run_config["histo"]["inc"]["max_mass"]),
|
||||
run_config["histo"]["inc"]["n_bins"],
|
||||
)
|
||||
inc_mass_bin_ctrs = bin_lims_to_ctrs(inc_mass_bin_lims)
|
||||
|
||||
scatt_bin_lims = np.logspace(
|
||||
np.log10(run_config["histo"]["scatt"]["min_D"]),
|
||||
np.log10(run_config["histo"]["scatt"]["max_D"]),
|
||||
run_config["histo"]["scatt"]["n_bins"],
|
||||
)
|
||||
scatt_bin_ctrs = bin_lims_to_ctrs(scatt_bin_lims)
|
||||
|
||||
timelag_bins_lims = np.linspace(
|
||||
run_config["histo"]["timelag"]["min"],
|
||||
run_config["histo"]["timelag"]["max"],
|
||||
run_config["histo"]["timelag"]["n_bins"],
|
||||
)
|
||||
timelag_bin_ctrs = bin_lims_to_ctrs(timelag_bins_lims)
|
||||
|
||||
for chunk_start, chunk_end in time_chunks:
|
||||
print(f"Processing: {chunk_start} to {chunk_end}")
|
||||
|
||||
pbp_filters = [
|
||||
("date", ">=", chunk_start.date().strftime("%Y-%m-%d")),
|
||||
("date", "<", chunk_end.date().strftime("%Y-%m-%d")),
|
||||
]
|
||||
if "hour" in run_config["chunking"]["freq"]: # optionally filter by hour
|
||||
pbp_filters.append(("hour", ">=", chunk_start.hour))
|
||||
pbp_filters.append(("hour", "<", chunk_end.hour))
|
||||
|
||||
# 2. HK processing --------------------------------------
|
||||
try:
|
||||
ddf_hk = dd.read_parquet(
|
||||
run_config["input_hk"],
|
||||
engine="pyarrow",
|
||||
filters=pbp_filters,
|
||||
calculate_divisions=True,
|
||||
)
|
||||
except (FileNotFoundError, OSError):
|
||||
print(" → no HK files for this chunk; skipping.")
|
||||
continue
|
||||
|
||||
if ddf_hk.npartitions == 0 or partition_rowcount(ddf_hk) == 0:
|
||||
print(" → HK frame is empty; skipping.")
|
||||
continue
|
||||
ddf_hk = ddf_hk.map_partitions(lambda pdf: pdf.sort_index())
|
||||
if not ddf_hk.known_divisions:
|
||||
ddf_hk = ddf_hk.reset_index().set_index( # 'calculated_time' becomes a column
|
||||
"calculated_time", sorted=False, shuffle="tasks"
|
||||
) # Dask now infers divisions
|
||||
ddf_hk = ddf_hk.repartition(freq="1h")
|
||||
meta = pd.DataFrame(
|
||||
{
|
||||
"Sample Flow Controller Read (sccm)": pd.Series(dtype="float64"),
|
||||
"Sample Flow Controller Read (vccm)": pd.Series(dtype="float64"),
|
||||
"date": pd.Series(dtype="datetime64[ns]"),
|
||||
"hour": pd.Series(dtype="int64"),
|
||||
},
|
||||
index=pd.DatetimeIndex([]),
|
||||
)
|
||||
ddf_hk = ddf_hk.repartition(freq="1h")
|
||||
meta = pd.DataFrame(
|
||||
{
|
||||
"Sample Flow Controller Read (sccm)": pd.Series(dtype="float64"),
|
||||
"Sample Flow Controller Read (vccm)": pd.Series(dtype="float64"),
|
||||
"date": pd.Series(dtype="datetime64[ns]"),
|
||||
"hour": pd.Series(dtype="int64"),
|
||||
},
|
||||
index=pd.DatetimeIndex([]),
|
||||
)
|
||||
ddf_hk_dt = ddf_hk.map_partitions(
|
||||
resample_hk_partition, dt=f"{run_config['dt']}s", meta=meta
|
||||
)
|
||||
|
||||
flow_dt = ddf_hk_dt["Sample Flow Controller Read (vccm)"].compute()
|
||||
|
||||
# 3. PBP processing --------------------------------------
|
||||
try:
|
||||
ddf_raw = dd.read_parquet(
|
||||
run_config["input_pbp"],
|
||||
engine="pyarrow",
|
||||
filters=pbp_filters,
|
||||
calculate_divisions=True,
|
||||
ddf_hk_dt = ddf_hk.map_partitions(
|
||||
resample_hk_partition, dt=f"{run_config['dt']}s", meta=meta
|
||||
)
|
||||
except (FileNotFoundError, OSError):
|
||||
print(" → no PbP files for this chunk; skipping.")
|
||||
continue
|
||||
|
||||
if ddf_raw.npartitions == 0 or partition_rowcount(ddf_raw) == 0:
|
||||
print(" → PbP frame is empty; skipping.")
|
||||
continue
|
||||
flow_dt = ddf_hk_dt["Sample Flow Controller Read (vccm)"].compute()
|
||||
|
||||
ddf_raw = ddf_raw.map_partitions(lambda pdf: pdf.sort_index())
|
||||
if not ddf_raw.known_divisions:
|
||||
ddf_raw = (
|
||||
ddf_raw.reset_index().set_index( # 'calculated_time' becomes a column
|
||||
# 3. PBP processing --------------------------------------
|
||||
try:
|
||||
ddf_raw = dd.read_parquet(
|
||||
run_config["input_pbp"],
|
||||
engine="pyarrow",
|
||||
filters=pbp_filters,
|
||||
calculate_divisions=True,
|
||||
)
|
||||
except (FileNotFoundError, OSError):
|
||||
print(" → no PbP files for this chunk; skipping.")
|
||||
continue
|
||||
|
||||
if ddf_raw.npartitions == 0 or partition_rowcount(ddf_raw) == 0:
|
||||
print(" → PbP frame is empty; skipping.")
|
||||
continue
|
||||
|
||||
ddf_raw = ddf_raw.map_partitions(lambda pdf: pdf.sort_index())
|
||||
if not ddf_raw.known_divisions:
|
||||
ddf_raw = ddf_raw.reset_index().set_index( # 'calculated_time' becomes a column
|
||||
"calculated_time", sorted=False, shuffle="tasks"
|
||||
) # Dask now infers divisions
|
||||
ddf_raw = ddf_raw.repartition(freq="1h")
|
||||
|
||||
ddf_cal = calibrate_single_particle(ddf_raw, instr_config, run_config)
|
||||
|
||||
ddf_pbp_with_flow = join_pbp_with_flow(ddf_cal, flow_dt, run_config)
|
||||
|
||||
delete_partition_if_exists(
|
||||
output_path=f"{run_config['output']}/pbp_calibrated",
|
||||
partition_values={
|
||||
"date": chunk_start.strftime("%Y-%m-%d 00:00:00"),
|
||||
"hour": chunk_start.hour,
|
||||
},
|
||||
)
|
||||
ddf_raw = ddf_raw.repartition(freq="1h")
|
||||
ddf_pbp_with_flow = enforce_schema(ddf_pbp_with_flow)
|
||||
|
||||
ddf_cal = calibrate_single_particle(ddf_raw, instr_config, run_config)
|
||||
|
||||
ddf_pbp_with_flow = join_pbp_with_flow(ddf_cal, flow_dt, run_config)
|
||||
|
||||
delete_partition_if_exists(
|
||||
output_path=f"{run_config['output']}/pbp_calibrated",
|
||||
partition_values={
|
||||
"date": chunk_start.strftime("%Y-%m-%d 00:00:00"),
|
||||
"hour": chunk_start.hour,
|
||||
},
|
||||
)
|
||||
ddf_pbp_with_flow = enforce_schema(ddf_pbp_with_flow)
|
||||
|
||||
ddf_pbp_with_flow.to_parquet(
|
||||
path=f"{run_config['output']}/pbp_calibrated",
|
||||
partition_on=["date", "hour"],
|
||||
engine="pyarrow",
|
||||
write_index=True,
|
||||
write_metadata_file=True,
|
||||
append=True,
|
||||
schema="infer",
|
||||
)
|
||||
|
||||
# 4. Aggregate PBP ---------------------------------------------
|
||||
ddf_pbp_dt = ddf_cal.map_partitions(
|
||||
build_dt_summary,
|
||||
dt_s=run_config["dt"],
|
||||
meta=build_dt_summary(ddf_cal._meta),
|
||||
)
|
||||
|
||||
ddf_pbp_hk_dt = aggregate_dt(ddf_pbp_dt, ddf_hk_dt, run_config)
|
||||
|
||||
# 4. (optional) dt bulk conc --------------------------
|
||||
if run_config["do_conc"]:
|
||||
meta_conc = add_concentrations(ddf_pbp_hk_dt._meta, dt=run_config["dt"])
|
||||
meta_conc = meta_conc.astype(
|
||||
{c: CANONICAL_DTYPES.get(c, DEFAULT_FLOAT) for c in meta_conc.columns},
|
||||
copy=False,
|
||||
).convert_dtypes(dtype_backend="pyarrow")
|
||||
|
||||
ddf_conc = ddf_pbp_hk_dt.map_partitions(
|
||||
add_concentrations, dt=run_config["dt"], meta=meta_conc
|
||||
).map_partitions(cast_and_arrow, meta=meta_conc)
|
||||
|
||||
idx_target = "datetime64[ns]"
|
||||
ddf_conc = ddf_conc.map_partitions(
|
||||
lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)),
|
||||
meta=ddf_conc._meta,
|
||||
)
|
||||
|
||||
# 2) cast partition columns *before* Dask strips them off
|
||||
ddf_conc["date"] = dd.to_datetime(ddf_conc["date"]).astype("datetime64[ns]")
|
||||
ddf_conc["hour"] = ddf_conc["hour"].astype("int64")
|
||||
|
||||
ddf_conc.to_parquet(
|
||||
f"{run_config['output']}/conc_{run_config['dt']}s",
|
||||
ddf_pbp_with_flow.to_parquet(
|
||||
path=f"{run_config['output']}/pbp_calibrated",
|
||||
partition_on=["date", "hour"],
|
||||
engine="pyarrow",
|
||||
write_index=True,
|
||||
@@ -212,172 +182,220 @@ def main():
|
||||
schema="infer",
|
||||
)
|
||||
|
||||
# 5. (optional) dt histograms --------------------------
|
||||
# 4. Aggregate PBP ---------------------------------------------
|
||||
ddf_pbp_dt = ddf_cal.map_partitions(
|
||||
build_dt_summary,
|
||||
dt_s=run_config["dt"],
|
||||
meta=build_dt_summary(ddf_cal._meta),
|
||||
)
|
||||
|
||||
if run_config["do_BC_hist"]:
|
||||
print("Computing BC distributions...")
|
||||
# --- Mass histogram
|
||||
BC_hist_configs = [
|
||||
{"flag_col": None, "flag_value": None},
|
||||
{"flag_col": "cnts_thin", "flag_value": 1},
|
||||
{"flag_col": "cnts_thin_noScatt", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_thin_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_ntl_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_ntl", "flag_value": 1},
|
||||
{
|
||||
"flag_col": "cnts_extreme_positive_timelag",
|
||||
"flag_value": 1,
|
||||
},
|
||||
{
|
||||
"flag_col": "cnts_thin_low_inc_scatt_ratio",
|
||||
"flag_value": 1,
|
||||
},
|
||||
{"flag_col": "cnts_thin_total", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick_total", "flag_value": 1},
|
||||
{"flag_col": "cnts_unclassified", "flag_value": 1},
|
||||
]
|
||||
ddf_pbp_hk_dt = aggregate_dt(ddf_pbp_dt, ddf_hk_dt, run_config)
|
||||
|
||||
results = []
|
||||
# 4. (optional) dt bulk conc --------------------------
|
||||
if run_config["do_conc"]:
|
||||
meta_conc = add_concentrations(ddf_pbp_hk_dt._meta, dt=run_config["dt"])
|
||||
meta_conc = meta_conc.astype(
|
||||
{
|
||||
c: CANONICAL_DTYPES.get(c, DEFAULT_FLOAT)
|
||||
for c in meta_conc.columns
|
||||
},
|
||||
copy=False,
|
||||
).convert_dtypes(dtype_backend="pyarrow")
|
||||
|
||||
for cfg_hist in BC_hist_configs[:2]:
|
||||
meta_hist = (
|
||||
make_hist_meta(
|
||||
bin_ctrs=inc_mass_bin_ctrs,
|
||||
kind="mass",
|
||||
ddf_conc = ddf_pbp_hk_dt.map_partitions(
|
||||
add_concentrations, dt=run_config["dt"], meta=meta_conc
|
||||
).map_partitions(cast_and_arrow, meta=meta_conc)
|
||||
|
||||
idx_target = "datetime64[ns]"
|
||||
ddf_conc = ddf_conc.map_partitions(
|
||||
lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)),
|
||||
meta=ddf_conc._meta,
|
||||
)
|
||||
|
||||
# 2) cast partition columns *before* Dask strips them off
|
||||
ddf_conc["date"] = dd.to_datetime(ddf_conc["date"]).astype(
|
||||
"datetime64[ns]"
|
||||
)
|
||||
ddf_conc["hour"] = ddf_conc["hour"].astype("int64")
|
||||
|
||||
ddf_conc.to_parquet(
|
||||
f"{run_config['output']}/conc_{run_config['dt']}s",
|
||||
partition_on=["date", "hour"],
|
||||
engine="pyarrow",
|
||||
write_index=True,
|
||||
write_metadata_file=True,
|
||||
append=True,
|
||||
schema="infer",
|
||||
)
|
||||
|
||||
# 5. (optional) dt histograms --------------------------
|
||||
|
||||
if run_config["do_BC_hist"]:
|
||||
print("Computing BC distributions...")
|
||||
# --- Mass histogram
|
||||
BC_hist_configs = [
|
||||
{"flag_col": None, "flag_value": None},
|
||||
{"flag_col": "cnts_thin", "flag_value": 1},
|
||||
{"flag_col": "cnts_thin_noScatt", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_thin_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_ntl_sat", "flag_value": 1},
|
||||
{"flag_col": "cnts_ntl", "flag_value": 1},
|
||||
{
|
||||
"flag_col": "cnts_extreme_positive_timelag",
|
||||
"flag_value": 1,
|
||||
},
|
||||
{
|
||||
"flag_col": "cnts_thin_low_inc_scatt_ratio",
|
||||
"flag_value": 1,
|
||||
},
|
||||
{"flag_col": "cnts_thin_total", "flag_value": 1},
|
||||
{"flag_col": "cnts_thick_total", "flag_value": 1},
|
||||
{"flag_col": "cnts_unclassified", "flag_value": 1},
|
||||
]
|
||||
|
||||
results = []
|
||||
|
||||
for cfg_hist in BC_hist_configs[:2]:
|
||||
meta_hist = (
|
||||
make_hist_meta(
|
||||
bin_ctrs=inc_mass_bin_ctrs,
|
||||
kind="mass",
|
||||
flag_col=cfg_hist["flag_col"],
|
||||
rho_eff=run_config["rho_eff"],
|
||||
BC_type=run_config["BC_type"],
|
||||
)
|
||||
.astype(DEFAULT_FLOAT, copy=False)
|
||||
.convert_dtypes(dtype_backend="pyarrow")
|
||||
)
|
||||
ddf_out = ddf_pbp_with_flow.map_partitions(
|
||||
process_hist_and_dist_partition,
|
||||
col="BC mass within range",
|
||||
flag_col=cfg_hist["flag_col"],
|
||||
flag_value=cfg_hist["flag_value"],
|
||||
bin_lims=inc_mass_bin_lims,
|
||||
bin_ctrs=inc_mass_bin_ctrs,
|
||||
dt=run_config["dt"],
|
||||
calculate_conc=True,
|
||||
flow=None,
|
||||
rho_eff=run_config["rho_eff"],
|
||||
BC_type=run_config["BC_type"],
|
||||
t=1,
|
||||
meta=meta_hist,
|
||||
).map_partitions(cast_and_arrow, meta=meta_hist)
|
||||
results.append(ddf_out)
|
||||
|
||||
# --- Scattering histogram
|
||||
if run_config["do_scatt_hist"]:
|
||||
print("Computing scattering distribution...")
|
||||
meta_hist = (
|
||||
make_hist_meta(
|
||||
bin_ctrs=scatt_bin_ctrs,
|
||||
kind="scatt",
|
||||
flag_col=None,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
)
|
||||
.astype(DEFAULT_FLOAT, copy=False)
|
||||
.convert_dtypes(dtype_backend="pyarrow")
|
||||
)
|
||||
ddf_out = ddf_pbp_with_flow.map_partitions(
|
||||
ddf_scatt = ddf_pbp_with_flow.map_partitions(
|
||||
process_hist_and_dist_partition,
|
||||
col="BC mass within range",
|
||||
flag_col=cfg_hist["flag_col"],
|
||||
flag_value=cfg_hist["flag_value"],
|
||||
bin_lims=inc_mass_bin_lims,
|
||||
bin_ctrs=inc_mass_bin_ctrs,
|
||||
col="Opt diam scatt only",
|
||||
flag_col=None,
|
||||
flag_value=None,
|
||||
bin_lims=scatt_bin_lims,
|
||||
bin_ctrs=scatt_bin_ctrs,
|
||||
dt=run_config["dt"],
|
||||
calculate_conc=True,
|
||||
flow=None,
|
||||
rho_eff=run_config["rho_eff"],
|
||||
BC_type=run_config["BC_type"],
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
t=1,
|
||||
meta=meta_hist,
|
||||
).map_partitions(cast_and_arrow, meta=meta_hist)
|
||||
results.append(ddf_out)
|
||||
results.append(ddf_scatt)
|
||||
|
||||
# --- Scattering histogram
|
||||
if run_config["do_scatt_hist"]:
|
||||
print("Computing scattering distribution...")
|
||||
meta_hist = (
|
||||
make_hist_meta(
|
||||
bin_ctrs=scatt_bin_ctrs,
|
||||
kind="scatt",
|
||||
flag_col=None,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
# --- Timelag histogram
|
||||
if run_config["do_timelag_hist"]:
|
||||
print("Computing time delay distribution...")
|
||||
mass_bins = (
|
||||
ddf_pbp_with_flow[["BC mass bin"]]
|
||||
.compute()
|
||||
.astype("Int64")
|
||||
.drop_duplicates()
|
||||
.dropna()
|
||||
)
|
||||
.astype(DEFAULT_FLOAT, copy=False)
|
||||
.convert_dtypes(dtype_backend="pyarrow")
|
||||
)
|
||||
ddf_scatt = ddf_pbp_with_flow.map_partitions(
|
||||
process_hist_and_dist_partition,
|
||||
col="Opt diam scatt only",
|
||||
flag_col=None,
|
||||
flag_value=None,
|
||||
bin_lims=scatt_bin_lims,
|
||||
bin_ctrs=scatt_bin_ctrs,
|
||||
dt=run_config["dt"],
|
||||
calculate_conc=True,
|
||||
flow=None,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
t=1,
|
||||
meta=meta_hist,
|
||||
).map_partitions(cast_and_arrow, meta=meta_hist)
|
||||
results.append(ddf_scatt)
|
||||
|
||||
# --- Timelag histogram
|
||||
if run_config["do_timelag_hist"]:
|
||||
print("Computing time delay distribution...")
|
||||
mass_bins = (
|
||||
ddf_pbp_with_flow[["BC mass bin"]]
|
||||
.compute()
|
||||
.astype("Int64")
|
||||
.drop_duplicates()
|
||||
.dropna()
|
||||
for idx, mass_bin in enumerate(mass_bins[:1]):
|
||||
ddf_bin = ddf_pbp_with_flow[
|
||||
ddf_pbp_with_flow["BC mass bin"] == mass_bin
|
||||
]
|
||||
|
||||
name_prefix = f"dNdlogDmev_{inc_mass_bin_ctrs[idx]:.2f}_timelag"
|
||||
|
||||
meta_hist = make_hist_meta(
|
||||
bin_ctrs=timelag_bin_ctrs,
|
||||
kind="timelag",
|
||||
flag_col="cnts_particles_for_tl_dist",
|
||||
name_prefix=name_prefix,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
)
|
||||
|
||||
tl_ddf = ddf_bin.map_partitions(
|
||||
process_hist_and_dist_partition,
|
||||
col="time_lag",
|
||||
flag_col="cnts_particles_for_tl_dist",
|
||||
flag_value=1,
|
||||
bin_lims=timelag_bins_lims,
|
||||
bin_ctrs=timelag_bin_ctrs,
|
||||
dt=run_config["dt"],
|
||||
calculate_conc=True,
|
||||
flow=None,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
t=1,
|
||||
name_prefix=name_prefix,
|
||||
meta=meta_hist,
|
||||
)
|
||||
|
||||
results.append(tl_ddf)
|
||||
|
||||
# --- Merge all hists
|
||||
merged_ddf = dd.concat(results, axis=1, interleave_partitions=True)
|
||||
|
||||
idx_target = "datetime64[ns]"
|
||||
merged_ddf = merged_ddf.map_partitions(
|
||||
lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)),
|
||||
meta=merged_ddf._meta,
|
||||
)
|
||||
|
||||
for idx, mass_bin in enumerate(mass_bins[:1]):
|
||||
ddf_bin = ddf_pbp_with_flow[
|
||||
ddf_pbp_with_flow["BC mass bin"] == mass_bin
|
||||
]
|
||||
index_as_dt = dd.to_datetime(merged_ddf.index.to_series())
|
||||
merged_ddf["date"] = index_as_dt.map_partitions(
|
||||
lambda s: s.dt.normalize(), meta=("date", "datetime64[ns]")
|
||||
)
|
||||
|
||||
name_prefix = f"dNdlogDmev_{inc_mass_bin_ctrs[idx]:.2f}_timelag"
|
||||
# --- Save hists to parquet
|
||||
|
||||
meta_hist = make_hist_meta(
|
||||
bin_ctrs=timelag_bin_ctrs,
|
||||
kind="timelag",
|
||||
flag_col="cnts_particles_for_tl_dist",
|
||||
name_prefix=name_prefix,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
)
|
||||
|
||||
tl_ddf = ddf_bin.map_partitions(
|
||||
process_hist_and_dist_partition,
|
||||
col="time_lag",
|
||||
flag_col="cnts_particles_for_tl_dist",
|
||||
flag_value=1,
|
||||
bin_lims=timelag_bins_lims,
|
||||
bin_ctrs=timelag_bin_ctrs,
|
||||
dt=run_config["dt"],
|
||||
calculate_conc=True,
|
||||
flow=None,
|
||||
rho_eff=None,
|
||||
BC_type=None,
|
||||
t=1,
|
||||
name_prefix=name_prefix,
|
||||
meta=meta_hist,
|
||||
)
|
||||
|
||||
results.append(tl_ddf)
|
||||
|
||||
# --- Merge all hists
|
||||
merged_ddf = dd.concat(results, axis=1, interleave_partitions=True)
|
||||
|
||||
idx_target = "datetime64[ns]"
|
||||
merged_ddf = merged_ddf.map_partitions(
|
||||
lambda pdf: pdf.set_index(pdf.index.astype(idx_target, copy=False)),
|
||||
meta=merged_ddf._meta,
|
||||
)
|
||||
|
||||
index_as_dt = dd.to_datetime(merged_ddf.index.to_series())
|
||||
merged_ddf["date"] = index_as_dt.map_partitions(
|
||||
lambda s: s.dt.normalize(), meta=("date", "datetime64[ns]")
|
||||
)
|
||||
|
||||
# --- Save hists to parquet
|
||||
|
||||
delete_partition_if_exists(
|
||||
output_path=f"{run_config['output']}/hists_{run_config['dt']}s",
|
||||
partition_values={
|
||||
"date": chunk_start.strftime("%Y-%m-%d"),
|
||||
"hour": chunk_start.hour,
|
||||
},
|
||||
)
|
||||
merged_ddf.to_parquet(
|
||||
f"{run_config['output']}/hists_{run_config['dt']}s",
|
||||
partition_on=["date"],
|
||||
append=True,
|
||||
schema="infer",
|
||||
)
|
||||
|
||||
client.close()
|
||||
delete_partition_if_exists(
|
||||
output_path=f"{run_config['output']}/hists_{run_config['dt']}s",
|
||||
partition_values={
|
||||
"date": chunk_start.strftime("%Y-%m-%d"),
|
||||
"hour": chunk_start.hour,
|
||||
},
|
||||
)
|
||||
merged_ddf.to_parquet(
|
||||
f"{run_config['output']}/hists_{run_config['dt']}s",
|
||||
partition_on=["date"],
|
||||
append=True,
|
||||
schema="infer",
|
||||
)
|
||||
finally:
|
||||
print("Final cleanup...", flush=True)
|
||||
client.close()
|
||||
cluster.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -114,7 +114,7 @@ def make_slurm_cluster(config):
|
||||
cluster.scale(1)
|
||||
client.wait_for_workers(1, timeout=600)
|
||||
print(f"Dask SLURM dashboard: {client.dashboard_link}")
|
||||
return client
|
||||
return client, cluster
|
||||
|
||||
|
||||
def make_local_cluster(config):
|
||||
|
||||
Reference in New Issue
Block a user