anemoi-datasets 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.4'
-__version_tuple__ = version_tuple = (0, 4, 4)
+__version__ = version = '0.5.0'
+__version_tuple__ = version_tuple = (0, 5, 0)

anemoi/datasets/commands/cleanup.py

@@ -0,0 +1,44 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Cleanup(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, subparser):
+        subparser.add_argument("path", help="Path to store the created data.")
+        subparser.add_argument(
+            "--delta",
+            help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
+            nargs="+",
+        )
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+        step = self.__class__.__name__.lower()
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+
+        task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = Cleanup

anemoi/datasets/commands/create.py

@@ -8,8 +8,6 @@ from concurrent.futures import as_completed
 import tqdm
 from anemoi.utils.humanize import seconds_to_human
 
-from anemoi.datasets.create.trace import enable_trace
-
 from . import Command
 
 LOG = logging.getLogger(__name__)
@@ -21,17 +19,16 @@ def task(what, options, *args, **kwargs):
     """
 
     now = datetime.datetime.now()
-    LOG.debug(f"Task {what}({args},{kwargs}) starting")
+    LOG.info(f"🎬 Task {what}({args},{kwargs}) starting")
 
-    from anemoi.datasets.create import Creator
+    from anemoi.datasets.create import creator_factory
 
-    if "trace" in options:
-        enable_trace(options["trace"])
+    options = {k: v for k, v in options.items() if v is not None}
 
-    c = Creator(**options)
-    result = getattr(c, what)(*args, **kwargs)
+    c = creator_factory(what.replace("-", "_"), **options)
+    result = c.run()
 
-    LOG.debug(f"Task {what}({args},{kwargs}) completed ({datetime.datetime.now()-now})")
+    LOG.info(f"🏁 Task {what}({args},{kwargs}) completed ({datetime.datetime.now()-now})")
     return result
 
 
@@ -60,6 +57,7 @@ class Create(Command):
         command_parser.add_argument("--trace", action="store_true")
 
     def run(self, args):
+
         now = time.time()
         if args.threads + args.processes:
             self.parallel_create(args)
@@ -68,11 +66,23 @@ class Create(Command):
         LOG.info(f"Create completed in {seconds_to_human(time.time()-now)}")
 
     def serial_create(self, args):
-        from anemoi.datasets.create import Creator
 
         options = vars(args)
-        c = Creator(**options)
-        c.create()
+        options.pop("command")
+        options.pop("threads")
+        options.pop("processes")
+
+        task("init", options)
+        task("load", options)
+        task("finalise", options)
+
+        task("patch", options)
+
+        task("init_additions", options)
+        task("run_additions", options)
+        task("finalise_additions", options)
+        task("cleanup", options)
+        task("verify", options)
 
     def parallel_create(self, args):
         """Some modules, like fsspec do not work well with fork()
@@ -82,10 +92,15 @@ class Create(Command):
         """
 
         options = vars(args)
-        parallel = args.threads + args.processes
-        args.use_threads = args.threads > 0
+        options.pop("command")
+
+        threads = options.pop("threads")
+        processes = options.pop("processes")
 
-        if args.use_threads:
+        use_threads = threads > 0
+        options["use_threads"] = use_threads
+
+        if use_threads:
             ExecutorClass = ThreadPoolExecutor
         else:
             ExecutorClass = ProcessPoolExecutor
@@ -95,18 +110,34 @@ class Create(Command):
 
         futures = []
 
+        parallel = threads + processes
+        with ExecutorClass(max_workers=parallel) as executor:
+            for n in range(total):
+                opt = options.copy()
+                opt["parts"] = f"{n+1}/{total}"
+                futures.append(executor.submit(task, "load", opt))
+
+            for future in tqdm.tqdm(as_completed(futures), desc="Loading", total=len(futures), colour="green", position=parallel + 1):  # fmt: skip
+                future.result()
+
+        with ExecutorClass(max_workers=1) as executor:
+            executor.submit(task, "finalise", options).result()
+
+        with ExecutorClass(max_workers=1) as executor:
+            executor.submit(task, "init-additions", options).result()
+
         with ExecutorClass(max_workers=parallel) as executor:
+            opt = options.copy()
+            opt["parts"] = f"{n+1}/{total}"
+            futures.append(executor.submit(task, "load", opt))
             for n in range(total):
-                futures.append(executor.submit(task, "load", options, parts=f"{n+1}/{total}"))
+                futures.append(executor.submit(task, "load-additions", opt))
 
-            for future in tqdm.tqdm(
-                as_completed(futures), desc="Loading", total=len(futures), colour="green", position=parallel + 1
-            ):
+            for future in tqdm.tqdm(as_completed(futures), desc="Computing additions", total=len(futures), colour="green", position=parallel + 1):  # fmt: skip
                 future.result()
 
         with ExecutorClass(max_workers=1) as executor:
-            executor.submit(task, "statistics", options).result()
-            executor.submit(task, "additions", options).result()
+            executor.submit(task, "finalise-additions", options).result()
             executor.submit(task, "cleanup", options).result()
             executor.submit(task, "verify", options).result()
 

anemoi/datasets/commands/finalise-additions.py

@@ -0,0 +1,45 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class FinaliseAdditions(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument(
+            "--delta",
+            help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
+            nargs="+",
+        )
+
+        command_parser.add_argument("path", help="Path to store the created data.")
+        command_parser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        step = "finalise-additions"
+        now = time.time()
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+            task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = FinaliseAdditions

anemoi/datasets/commands/finalise.py

@@ -0,0 +1,39 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Finalise(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("path", help="Path to store the created data.")
+        command_parser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+        step = "finalise"
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+            task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = Finalise

anemoi/datasets/commands/init-additions.py

@@ -0,0 +1,45 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class InitAdditions(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument(
+            "--delta",
+            help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
+            nargs="+",
+        )
+
+        command_parser.add_argument("path", help="Path to store the created data.")
+        command_parser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        step = "init-additions"
+        now = time.time()
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+            task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = InitAdditions

anemoi/datasets/commands/init.py

@@ -0,0 +1,67 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Init(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, subparser):
+
+        subparser.add_argument("config", help="Configuration yaml file defining the recipe to create the dataset.")
+        subparser.add_argument("path", help="Path to store the created data.")
+
+        subparser.add_argument(
+            "--overwrite",
+            action="store_true",
+            help="Overwrite existing files. This will delete the target dataset if it already exists.",
+        )
+        subparser.add_argument(
+            "--test",
+            action="store_true",
+            help="Build a small dataset, using only the first dates. And, when possible, using low resolution and less ensemble members.",
+        )
+        subparser.add_argument(
+            "--check-name",
+            dest="check_name",
+            action="store_true",
+            help="Check if the dataset name is valid before creating it.",
+        )
+        subparser.add_argument(
+            "--no-check-name",
+            dest="check_name",
+            action="store_false",
+            help="Do not check if the dataset name is valid before creating it.",
+        )
+        subparser.set_defaults(check_name=False)
+        subparser.add_argument("--cache", help="Location to store the downloaded data.", metavar="DIR")
+
+        subparser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+
+        task("init", options)
+
+        LOG.info(f"Create step '{self.__class__.__name__.lower()}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = Init

anemoi/datasets/commands/inspect.py

@@ -149,7 +149,7 @@ class Version:
         print()
         print(f'📅 Start      : {self.first_date.strftime("%Y-%m-%d %H:%M")}')
         print(f'📅 End        : {self.last_date.strftime("%Y-%m-%d %H:%M")}')
-        print(f"⏰ Frequency  : {self.frequency}h")
+        print(f"⏰ Frequency  : {self.frequency}")
         if self.n_missing_dates is not None:
             print(f"🚫 Missing    : {self.n_missing_dates:,}")
         print(f"🌎 Resolution : {self.resolution}")

anemoi/datasets/commands/load-additions.py

@@ -0,0 +1,47 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class LoadAdditions(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, subparser):
+
+        subparser.add_argument("--parts", nargs="+", help="Only load the specified parts of the dataset.")
+        subparser.add_argument(
+            "--delta",
+            help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
+            nargs="+",
+        )
+
+        subparser.add_argument("path", help="Path to store the created data.")
+        subparser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+        step = "load-additions"
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+            task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = LoadAdditions

anemoi/datasets/commands/load.py

@@ -0,0 +1,47 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Load(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, subparser):
+
+        subparser.add_argument("--parts", nargs="+", help="Only load the specified parts of the dataset.")
+        # subparser.add_argument(
+        #        "--delta",
+        #        help="Compute statistics tendencies on a given time delta, if possible. Must be a multiple of the frequency.",
+        #    )
+
+        subparser.add_argument("path", help="Path to store the created data.")
+        subparser.add_argument("--cache", help="Location to store the downloaded data.", metavar="DIR")
+        subparser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+        step = "load"
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+            task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = Load

anemoi/datasets/commands/patch.py

@@ -0,0 +1,39 @@
+import logging
+import time
+
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.commands.create import task
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Patch(Command):
+    """Create a dataset, step by step."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, parser):
+        parser.add_argument("path", help="Path to store the created data.")
+
+    def run(self, args):
+        options = vars(args)
+        options.pop("command")
+        now = time.time()
+        step = self.__class__.__name__.lower()
+
+        if "version" in options:
+            options.pop("version")
+
+        if "debug" in options:
+            options.pop("debug")
+
+        task(step, options)
+
+        LOG.info(f"Create step '{step}' completed in {seconds_to_human(time.time()-now)}")
+
+
+command = Patch