Survey IO

Read and write survey files, the self-contained HDF5 output of hifuku survey. The --description field takes a free-form note; markdown is recommended.

save_survey

save_survey(path, *, taxon_table, triangle, gene_labels, alignment_paths, results, model_label, device, seed, n_iters, description='', author=None) -> None

Write a survey file: namespace, genes, anchor graph, one chart, the per-gene archives (with elite newick trees), and provenance.

results is one ArchiveResult per gene, aligned with gene_labels.

Source code in src/hifuku/survey_io.py

def save_survey(path, *, taxon_table, triangle, gene_labels, alignment_paths,
                results, model_label, device, seed, n_iters,
                description="", author=None) -> None:
    """Write a survey file: namespace, genes, anchor graph, one chart, the
    per-gene archives (with elite newick trees), and provenance.

    ``results`` is one ``ArchiveResult`` per gene, aligned with ``gene_labels``.
    """
    n_genes = len(gene_labels)
    m = triangle.tree_metric
    with h5py.File(path, "w") as f:
        f.attrs["schema_version"] = _SCHEMA_VERSION
        f.attrs["field_type"] = _FIELD_TYPE
        f.attrs["metric"] = "normalized_branch_score"
        f.attrs["seed"] = seed
        f.attrs["n_genes"] = n_genes
        f.attrs["n_charts"] = 1
        f.attrs["created"] = datetime.datetime.now(datetime.timezone.utc).isoformat()
        f.attrs["description"] = description

        f.create_dataset("namespace/taxa",
                         data=np.array(taxon_table.labels, dtype=object),
                         dtype=_STR_DTYPE)

        anchors = (triangle.r1, triangle.r2, triangle.r3)
        for i, (label, apath) in enumerate(zip(gene_labels, alignment_paths)):
            g = f.create_group(f"genes/{i}")
            g.attrs["label"] = label
            g.attrs["alignment_path"] = str(apath) if apath else ""
            g.attrs["alignment_sha256"] = _sha256(apath) if apath else ""
            g.create_dataset("nj_newick", data=anchors[i].as_newick(),
                             dtype=_STR_DTYPE)

        ag = f.create_group("anchor_graph")
        ag.attrs["anchor_gene_indices"] = np.arange(3, dtype=np.int32)
        ag.create_dataset("dist", data=triangle.dist.astype(np.float64))
        ag.create_dataset("qual", data=triangle.qual.astype(np.float64))

        c = f.create_group("charts/0")
        c.attrs["anchor_gene_indices"] = np.arange(3, dtype=np.int32)
        c.attrs["metric_w"] = float(m.w)
        c.attrs["c_ref"] = float(m.c_ref)
        c.attrs["r_ref"] = float(m.r_ref)
        c.attrs["cell"] = float(results[0].archive.cell)
        c.attrs["origin"] = float(results[0].archive.origin)
        from hifuku.anchor_triangle import triangle_quality
        Q, _, _ = triangle_quality(triangle.P1, triangle.P2, triangle.P3)
        c.attrs["Q"] = float(Q)
        c.create_dataset("P", data=np.stack(
            [triangle.P1, triangle.P2, triangle.P3]).astype(np.float64))

        for k, res in enumerate(results):
            arch = res.archive
            gk = c.create_group(f"genes/{k}")
            keys = list(arch.keys)
            gk.attrs["gene_index"] = k
            gk.attrs["n_filled"] = arch.n_filled
            gk.attrs["best_logL"] = arch.best_logL()
            gk.attrs["converged"] = int(res.converged)
            gk.attrs["halt_iter"] = int(res.halt_iter)
            gk.attrs["device"] = device
            gk.attrs["model"] = model_label
            gk.attrs["n_iters"] = int(n_iters)
            gk.attrs["seed"] = int(seed)
            niche = np.array(keys, dtype=np.int32).reshape(-1, 2)
            logL = np.array([arch.elites[key].logL for key in keys], dtype=np.float64)
            z = np.array([arch.elites[key].z for key in keys], dtype=np.float64)
            newick = np.array([arch.tree_at(key).as_newick() for key in keys],
                              dtype=object)
            gk.create_dataset("niche_ij", data=niche)
            gk.create_dataset("logL", data=logL)
            gk.create_dataset("z", data=z)
            gk.create_dataset("tree", data=newick, dtype=_STR_DTYPE,
                              compression="gzip")
            gk.create_dataset("history", data=res.history.astype(np.float64))

        prov = f.create_group("provenance")
        for key, val in _provenance(device, author).items():
            prov.attrs[key] = val
        rows = [[str(p) if p else "", _sha256(p) if p else ""]
                for p in alignment_paths]
        prov.create_dataset("inputs", data=np.array(rows, dtype=object),
                            dtype=_STR_DTYPE)

load_survey

load_survey(path) -> Survey

Reconstruct a Survey from a survey HDF5 file (schema v1).

Source code in src/hifuku/survey_io.py

def load_survey(path) -> Survey:
    """Reconstruct a Survey from a survey HDF5 file (schema v1)."""
    from hifuku.taxa import TaxonTable
    from hifuku.tree_io import read_tree
    from hifuku.anchor_triangle import AnchorTriangle, CBSMetric
    from hifuku.archive import Elite, EliteArchive, _set_global_leaf_indices

    with h5py.File(path, "r") as f:
        if int(f.attrs.get("schema_version", 0)) != _SCHEMA_VERSION:
            raise ValueError(f"unsupported schema_version in {path}")
        description = _s(f.attrs.get("description", ""))
        provenance = {k: _s(v) for k, v in f["provenance"].attrs.items()}

        taxa_labels = f["namespace/taxa"][:]
        n_taxa = len(taxa_labels)
        tt = TaxonTable()
        for label in taxa_labels:
            tt.require(_s(label))

        genes = []
        gene_labels = []
        for i in range(int(f.attrs["n_genes"])):
            g = f[f"genes/{i}"]
            nj = read_tree(_s(g["nj_newick"][()]), tt, from_string=True)
            _set_global_leaf_indices(nj, tt)
            label = _s(g.attrs["label"])
            gene_labels.append(label)
            genes.append(GeneInfo(label, nj, _s(g.attrs["alignment_path"]),
                                  _s(g.attrs["alignment_sha256"])))

        charts = []
        for ci in range(int(f.attrs["n_charts"])):
            c = f[f"charts/{ci}"]
            idx = np.asarray(c.attrs["anchor_gene_indices"])
            r = [genes[int(j)].nj_tree for j in idx]
            metric = CBSMetric(w=float(c.attrs["metric_w"]),
                               c_ref=float(c.attrs["c_ref"]),
                               r_ref=float(c.attrs["r_ref"]))
            P = c["P"][:]
            tri = AnchorTriangle(P1=P[0], P2=P[1], P3=P[2],
                                 r1=r[0], r2=r[1], r3=r[2],
                                 dist=f["anchor_graph/dist"][:],
                                 qual=f["anchor_graph/qual"][:],
                                 metric=metric)
            cell = float(c.attrs["cell"])
            origin = float(c.attrs["origin"])
            archives = []
            meta = []
            for k in range(len(gene_labels)):
                gk = c[f"genes/{k}"]
                niche = gk["niche_ij"][:]
                logL = gk["logL"][:]
                z = gk["z"][:]
                newick = gk["tree"][:]
                src = {}
                arch = EliteArchive(cell=cell, origin=origin, anchors=tri,
                                    taxon_table=tt, metric=metric, tree_source=src)
                for row in range(niche.shape[0]):
                    key = (int(niche[row, 0]), int(niche[row, 1]))
                    arch.elites[key] = Elite(float(logL[row]), None, float(z[row]))
                    arch.keys.append(key)
                    src[key] = _s(newick[row])
                archives.append(arch)
                meta.append({"converged": bool(gk.attrs["converged"]),
                            "halt_iter": int(gk.attrs["halt_iter"])})
            charts.append(Chart(tri, archives, gene_labels, meta))

    return Survey(description, provenance, genes, charts, n_taxa)