Untitled

def compute_density(files: list, temp_var: str, salinity_var: str, output_var: str, output_dir: str = "."):
    """
    Compute the density by using the TEOS-10 package
    """
    pattern = r'\d{10}'
    match = re.search(pattern, files[0])
    date = match.group(0)

    _data = read_mfnetcdf(files, chunks={"time": "auto"})

    # set the dataset to only the variables we need
    _data = _data[[temp_var, salinity_var]]

    pressure = gsw.p_from_z(_data["depth"]*-1, _data["lat"])
    absolute_salinity = gsw.SA_from_SP(_data[salinity_var], pressure, _data["lon"], _data["lat"])
    conservative_temp = gsw.CT_from_t(absolute_salinity, _data[temp_var], pressure)
    density = gsw.rho(absolute_salinity, conservative_temp, pressure)

    # Update the attributes for the new dataarray
    density = density.assign_attrs(long_name="Density", units="kg/m")
    density = density.astype(np.float32)

    # assign density to the xarray dataset for output
    _data[output_var] = density

    # set the output dataset to the only output variable needed
    output_data = _data[[output_var]]

    times, datasets = zip(*output_data.groupby("time"))
    paths = [f"{output_dir}/density_{date}_{int(idx):04d}.nc" for idx, _ in enumerate(times)]
    xr.save_mfdataset(datasets, paths)
    _data.close()