zephyr-backend/modules/saber/__init__.py

import glob
from io import BytesIO
from quart import Blueprint, request, send_file
from matplotlib import pyplot as plt
from CONSTANT import DATA_BASEPATH
from modules.saber.planetw_perday import saber_planetw_plot_perday
from modules.saber.planetw_plot import saber_planetw_plot
from modules.saber.planetw_process import saber_planetw_process
from modules.saber.gravityw_process import SaberGravitywProcessor
from modules.saber.gravityw_render import SaberGravitywRenderer
from modules.saber.utils import *
from quart.utils import run_sync


saber_module = Blueprint('saber', __name__)
# lat_range = (30.0, 40.0)
# alt_range = (20.0, 105.0)
# lambda_range = (2, 15)
# lvboin = True

# processor = DataProcessor(
#     lat_range=lat_range, alt_range=alt_range, lambda_range=lambda_range, lvboin=lvboin)
renderer = SaberGravitywRenderer()


def get_gravityw_processer():
    lat_str = request.args.get("lat_range")
    # if found ~ in str, replace it with ,
    if "~" in lat_str:
        lat_str = lat_str.replace("~", ",")
    if lat_str is None:
        lat_str = "30.0,40.0"
    print(lat_str)
    lat_range = tuple(map(float, lat_str.split(",")))
    print(lat_range)
    alt_str = request.args.get("alt_range")
    if alt_str is None:
        alt_str = "20.0,105.0"
    alt_range = tuple(map(float, alt_str.split(",")))
    lambda_str = request.args.get("lambda_range")
    if lambda_str is None:
        lambda_str = "2,15"
    lambda_range = tuple(map(int, lambda_str.split(",")))
    lvboin = request.args.get("lvboin") == "true"
    if lvboin is None:
        lvboin = True

    _p = SaberGravitywProcessor(
        lat_range=lat_range, alt_range=alt_range, lambda_range=lambda_range, lvboin=lvboin)
    return _p


async def extract_payload():
    buffer = BytesIO()
    plt.savefig(buffer, format="png")
    buffer.seek(0)
    return await send_file(buffer, mimetype="image/png")


_all_saber_files = glob.glob(f"{DATA_BASEPATH.saber}/**/**.nc", recursive=True)


@saber_module.route("/metadata")
def get_files():
    # normalizing the path, and replace \\ with /
    all_saber_files = [path.replace("\\", "/") for path in _all_saber_files]
    return all_saber_files


@saber_module.route("/metadata/list_days")
async def get_days():
    path = request.args.get("path")

    ncfile = await run_sync(data_nc_load)(path)

    return ncfile.date_time.tolist()


@saber_module.route("/render/gravity_wave/per_day/wave_fitting")
@saber_module.route("/render/plot_wave_fitting")
async def do_plot_wave_day_fitting():
    path = request.args.get("path")
    day = request.args.get("day")
    height = request.args.get("height_no")

    ncfile = await run_sync(data_nc_load)(path)
    processor = await run_sync(get_gravityw_processer)()
    data = await run_sync(processor.process_day)(ncfile, int(day))
    await run_sync(renderer.plot_wave_fitting)(data, int(height))
    return await extract_payload()


@saber_module.route("/render/gravity_wave/per_day/fft_ifft")
@saber_module.route("/render/day_fft_ifft_plot")
async def do_day_fft_ifft_plot():
    path = request.args.get("path")
    day = request.args.get("day")
    cycle_no = request.args.get("cycle_no")

    ncfile = await run_sync(data_nc_load)(path)
    processor = await run_sync(get_gravityw_processer)()
    data = await run_sync(processor.process_day)(ncfile, int(day))
    # 高度滤波处理
    await run_sync(renderer.day_fft_ifft_plot)(wave_data=data, cycle_no=int(cycle_no))
    return await extract_payload()


@saber_module.route("/render/gravity_wave/per_day/power_wave_plot")
@saber_module.route("/render/day_cycle_power_wave_plot")
async def do_day_cycle_power_wave_plot():
    path = request.args.get("path")
    day = request.args.get("day")
    cycle_no = request.args.get("cycle_no")

    ncfile = await run_sync(data_nc_load)(path)
    processor = await run_sync(get_gravityw_processer)()
    data = await run_sync(processor.process_day)(ncfile, int(day))
    await run_sync(renderer.day_cycle_power_wave_plot)(wave_data=data, cycle_no=int(cycle_no))
    return await extract_payload()


@saber_module.route("/render/gravity_wave/per_month/power_wave_plot")
@saber_module.route("/render/month_power_wave_plot")
async def do_month_power_wave_plot():
    path = request.args.get("path")
    month = request.args.get("month")

    ncfile = await run_sync(data_nc_load)(path)
    processor = await run_sync(get_gravityw_processer)()
    data = await run_sync(processor.process_month)(ncfile)
    await run_sync(renderer.month_power_wave_plot)(wave_data=data, date_time=ncfile.date_time)
    return await extract_payload()


@saber_module.route("/render/gravity_wave/per_year/power_wave_plot")
@saber_module.route("/render/year_power_wave_plot")
async def do_year_power_wave_plot():
    year = request.args.get("year")
    processor = await run_sync(get_gravityw_processer)()
    data = await run_sync(processor.process_year)([
        data_nc_load(path) for path in glob.glob(f"{DATA_BASEPATH.saber}/{year}/*.nc")
    ])
    await run_sync(renderer.year_power_wave_plot)(year_wave=data)
    return await extract_payload()


@saber_module.route("/render/planet_wave/per_year/energy_plot")
# above path is named correctly
# below path is named wrongly, but kept due to compatibility
@saber_module.route("/render/gravity_year")
async def do_gravity_year():
    year = request.args.get("year")
    T = request.args.get("T")
    k = request.args.get("k")
    H = request.args.get("H")
    lat_target = request.args.get("lat_target")
    year = int(year)
    T = int(T)
    if T not in [5, 10, 16]:
        raise ValueError("T must be in [5, 10, 16]")
    data = await run_sync(saber_planetw_process)(
        year,
        target_h=int(H),
        target_lat=int(lat_target)
    )
    await run_sync(saber_planetw_plot)(data, T, k=int(k))
    return await extract_payload()


@saber_module.route("/render/planet_wave/per_day")
async def do_pl_year():
    year = request.args.get("year")
    T = request.args.get("T")
    # k = request.args.get("k")
    H = request.args.get("H")
    lat_target = request.args.get("lat_target")
    start_day = request.args.get("start_day", 0)
    year = int(year)
    T = int(T)
    if T not in [5, 10, 16]:
        raise ValueError("T must be in [5, 10, 16]")
    data = await run_sync(saber_planetw_process)(
        year,
        target_h=int(H),
        target_lat=int(lat_target)
    )
    await run_sync(saber_planetw_plot_perday)(data, T, int(start_day))
    return await extract_payload()