data_constants_3d.py

from source.data_source import *


class DataConstants3D(DataSource):

    def download(self, start_date: datetime, end_date: datetime) -> bool:
        """Generate constants for liquid border 3D.

        :param Source self: mock source
        :param datetime start_date: start date
        :param datetime end_date: end date
        :returns: bool
        """
        real_start_date, real_end_date, mode = self.compute_real_dates(start_date, end_date)
        duration_seconds = (real_end_date - real_start_date).total_seconds()
        nr_time_steps = int(duration_seconds // self.constants_time_step) + 1

        # organic load and sediments for ocean are set to zero
        self.variables[0].values = np.zeros(nr_time_steps)
        self.variables[1].values = np.zeros(nr_time_steps)
        self.variables[2].values = np.zeros(nr_time_steps)

        # For garonne and dordogne:
        # set organic load to 3.0
        # set nh4 to 0.054
        self.variables[3].values = 3.0 * np.ones(nr_time_steps)
        self.variables[4].values = 0.054 * np.ones(nr_time_steps)
        self.variables[5].values = 3.0 * np.ones(nr_time_steps)
        self.variables[6].values = 0.054 * np.ones(nr_time_steps)

        timestep_seconds = timedelta(seconds=self.constants_time_step)
        times: np.ndarray = np.arange(real_start_date, real_end_date + timestep_seconds, timestep_seconds)
        for idx in range(len(self.variables)):
            self.variables[idx].times = times

        return True