moradology/wtf.py Secret

## 76 changes: 76 additions & 0 deletions wtf.py
@@ -0,0 +1,76 @@

    #!/usr/bin/env python
#!/usr/bin/env python

    import json
import json

    import multiprocessing
import multiprocessing

    from typing import Dict, List, Optional, Sequence, Tuple
from typing import Dict, List, Optional, Sequence, Tuple


    from kerchunk.combine import MultiZarrToZarr
from kerchunk.combine import MultiZarrToZarr


    CONCAT_DIMS = ['time']
CONCAT_DIMS = ['time']

    IDENTICAL_DIMS = ['lat', 'lon']
IDENTICAL_DIMS = ['lat', 'lon']


    def load_refs(ldjson_file: str) -> list[dict]:
def load_refs(ldjson_file: str) -> list[dict]:

        refs = []
    refs = []

        with open(ldjson_file, "r") as f:
    with open(ldjson_file, "r") as f:

            for line in f:
        for line in f:

                refs.append(json.loads(line)[0])
            refs.append(json.loads(line)[0])

        return refs
    return refs


    def mzz(refs):
def mzz(refs):

        return MultiZarrToZarr(
    return MultiZarrToZarr(

            refs,
        refs,

            concat_dims=CONCAT_DIMS,
        concat_dims=CONCAT_DIMS,

            identical_dims=IDENTICAL_DIMS,
        identical_dims=IDENTICAL_DIMS,

            target_options={"anon": True},
        target_options={"anon": True},

            remote_options={"anon": True},
        remote_options={"anon": True},

            remote_protocol=None
        remote_protocol=None

        )
    )


    def merge_refs(refs: list[dict]) -> dict:
def merge_refs(refs: list[dict]) -> dict:

        return mzz(refs).translate()
    return mzz(refs).translate()


    # Distributed workflow
# Distributed workflow

    def worker_func(refs: list[dict]) -> MultiZarrToZarr:
def worker_func(refs: list[dict]) -> MultiZarrToZarr:

        def create_accumulator():
    def create_accumulator():

            return None
        return None

        def add_input(accumulator: MultiZarrToZarr, item: dict) -> MultiZarrToZarr:
    def add_input(accumulator: MultiZarrToZarr, item: dict) -> MultiZarrToZarr:

            if not accumulator:
        if not accumulator:

                references = [item]
            references = [item]

            else:
        else:

                references = [accumulator.translate(), item]
            references = [accumulator.translate(), item]

            return mzz(references)
        return mzz(references)

        acc = create_accumulator()
    acc = create_accumulator()

        for ref in refs:
    for ref in refs:

            acc = add_input(acc, ref)
        acc = add_input(acc, ref)

        return acc
    return acc


    def distributed_merge(refs: list[list[dict]]) -> dict:
def distributed_merge(refs: list[list[dict]]) -> dict:

        def merge_accumulators(accumulators: Sequence[MultiZarrToZarr]) -> MultiZarrToZarr:
    def merge_accumulators(accumulators: Sequence[MultiZarrToZarr]) -> MultiZarrToZarr:

            references = [a.translate() for a in accumulators]
        references = [a.translate() for a in accumulators]

            return mzz(references)
        return mzz(references)


        def extract_output(accumulator: MultiZarrToZarr) -> dict:
    def extract_output(accumulator: MultiZarrToZarr) -> dict:

            return accumulator.translate(),
        return accumulator.translate(),


        with multiprocessing.Pool(4) as p:
    with multiprocessing.Pool(4) as p:

            accumulators: list[MultiZarrToZarr] = p.map(worker_func, refs)
        accumulators: list[MultiZarrToZarr] = p.map(worker_func, refs)

        merged = merge_accumulators(accumulators)
    merged = merge_accumulators(accumulators)

        return extract_output(merged)
    return extract_output(merged)


    def compare_merge_size(single_dict, multi_dict):
def compare_merge_size(single_dict, multi_dict):

        single_bytes = len(json.dumps(single_dict).encode("utf-8"))
    single_bytes = len(json.dumps(single_dict).encode("utf-8"))

        multi_bytes = len(json.dumps(multi_dict).encode("utf-8"))
    multi_bytes = len(json.dumps(multi_dict).encode("utf-8"))

        print(f"The single process dict is {single_bytes}")
    print(f"The single process dict is {single_bytes}")

        print(f"The multi process dict is {multi_bytes}")
    print(f"The multi process dict is {multi_bytes}")


    def main():
def main():

        refs = load_refs("single/inputs_raw_15286.json")
    refs = load_refs("single/inputs_raw_15286.json")


        # Expected merge results
    # Expected merge results

        single_merge = merge_refs(refs)
    single_merge = merge_refs(refs)

        multi_refs = [[refs[0], refs[1]], [refs[2]], [refs[3]], [refs[4]]]
    multi_refs = [[refs[0], refs[1]], [refs[2]], [refs[3]], [refs[4]]]

        multi_merge = distributed_merge(multi_refs)
    multi_merge = distributed_merge(multi_refs)


        compare_merge_size(single_merge, multi_merge)
    compare_merge_size(single_merge, multi_merge)


    if __name__ == "__main__":
if __name__ == "__main__":

        main()
    main()