from __future__ import absolute_import, division
# from io import StringIO
from glob import glob
import os.path as op
import tempfile
import simplejson as json
from pandas.api import types
import numpy as np
import pandas as pd
# from _common import cooler_cmp
from click.testing import CliRunner
import cooler
import pytest
### INGEST AND AGGREGATION ###
from cooler.cli.cload import pairs as cload_pairs
from cooler.cli.load import load
tmp = tempfile.gettempdir()
testdir = op.realpath(op.dirname(__file__))
datadir = op.join(testdir, "data")
def _run_cload_pairs(runner, binsize, extra_args):
args = [
op.join(datadir, "toy.chrom.sizes") + ":" + str(binsize),
op.join(datadir, "toy.pairs"),
"toy.{}.cool".format(binsize),
"-c1", "2",
"-p1", "3",
"-c2", "4",
"-p2", "5",
"--assembly", "toy",
"--chunksize", "10",
] + extra_args
return runner.invoke(cload_pairs, args)
def _cmp_pixels_2_bg(f_out, f_ref, one_based_ref=True):
# output, 1-based starts
out_df = cooler.Cooler(f_out).pixels(join=True)[:]
if one_based_ref:
out_df["start1"] += 1
out_df["start2"] += 1
# reference
ref_df = pd.read_csv(
f_ref,
sep="\t",
names=["chrom1", "start1", "end1", "chrom2", "start2", "end2", "count"],
)
assert np.all(out_df == ref_df)
# '--no-symmetric-upper'
# '--input-copy-status', 'unique|duplex',
@pytest.mark.parametrize(
"ref,extra_args",
[
("symm.upper", []), # reflect triu pairs
("symm.upper", ["--input-copy-status", "unique"]), # reflect triu pairs
("asymm", ["--no-symmetric-upper"]),
],
)
def test_cload_symm_asymm(ref, extra_args):
runner = CliRunner()
with runner.isolated_filesystem():
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
_cmp_pixels_2_bg("toy.2.cool", op.join(datadir, "toy.{}.2.bg2".format(ref)))
# '--temp-dir', '',
# '--no-delete-temp',
# '--max-merge', '',
@pytest.mark.parametrize(
"ref,extra_args", [("symm.upper", ["--temp-dir", ".", "--no-delete-temp"])]
)
def test_cload_mergepass(ref, extra_args):
runner = CliRunner()
with runner.isolated_filesystem():
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
_cmp_pixels_2_bg("toy.2.cool", op.join(datadir, "toy.{}.2.bg2".format(ref)))
assert len(cooler.fileops.list_coolers(glob("*.cool")[0])) > 0
# '--field', '',
# '--no-count', '',
def test_cload_field():
runner = CliRunner()
with runner.isolated_filesystem():
extra_args = ["--field", "score=8"]
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
pixels = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels.columns and types.is_integer_dtype(
pixels.dtypes["count"]
)
assert "score" in pixels.columns and types.is_float_dtype(
pixels.dtypes["score"]
)
extra_args = ["--field", "count=8"]
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
pixels = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels.columns and types.is_integer_dtype(
pixels.dtypes["count"]
)
assert np.allclose(pixels["count"][:], 0)
extra_args = ["--field", "count=8:dtype=float"]
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
pixels = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels.columns and types.is_float_dtype(
pixels.dtypes["count"]
)
assert np.allclose(pixels["count"][:], 0.2)
extra_args = ["--field", "count=8:agg=min,dtype=float"]
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
pixels = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels.columns and types.is_float_dtype(
pixels.dtypes["count"]
)
assert np.allclose(pixels["count"][:], 0.1)
## don't implement the --no-count for now
# extra_args = ['--field', 'score=7:dtype=float', '--no-count']
# result = _run_cload_pairs(runner, 2, extra_args)
# assert result.exit_code == 0
# pixels = cooler.Cooler('toy.2.cool').pixels()[:]
# assert 'count' not in pixels.columns
# assert 'score' in pixels.columns and types.is_float_dtype(pixels.dtypes['score'])
# '--metadata', '',
# '--zero-based',
# '--comment-char', '',
# '--storage-options', '',
def test_cload_other_options():
runner = CliRunner()
with runner.isolated_filesystem():
meta = {"foo": "bar", "number": 42}
with open("meta.json", "w") as f:
json.dump(meta, f)
extra_args = [
"--metadata",
"meta.json",
"--zero-based",
"--storage-options",
"shuffle=True,fletcher32=True,compression=lzf",
]
result = _run_cload_pairs(runner, 2, extra_args)
assert result.exit_code == 0
c = cooler.Cooler("toy.2.cool")
assert c.info["metadata"] == meta
with c.open("r") as h5:
dset = h5["bins/start"]
assert dset.shuffle
assert dset.fletcher32
assert dset.compression == "lzf"
def _run_load(runner, matrix_file, format, binsize, extra_args):
args = [
"-f",
format,
op.join(datadir, "toy.chrom.sizes") + ":" + str(binsize),
op.join(datadir, matrix_file),
"toy.{}.cool".format(binsize),
"--assembly",
"toy",
"--chunksize",
"10",
] + extra_args
return runner.invoke(load, args)
# '--no-symmetric-upper'
# '--input-copy-status', 'unique|duplex',
@pytest.mark.parametrize(
"ref,extra_args",
[
("symm.upper", []), # reflect tril pairs
("symm.upper", ["--one-based", "--input-copy-status", "unique"]), # reflect tril pairs
("asymm", ["--one-based", "--no-symmetric-upper"]),
],
)
def test_load_symm_asymm(ref, extra_args):
runner = CliRunner()
with runner.isolated_filesystem():
ref = op.join(datadir, "toy.{}.2.bg2".format(ref))
result = _run_load(runner, ref, "bg2", 2, extra_args)
assert result.exit_code == 0
_cmp_pixels_2_bg("toy.2.cool", ref)
# '--field', '',
def test_load_field():
runner = CliRunner()
with runner.isolated_filesystem():
extra_args = ["--field", "count=7:dtype=float"]
result = _run_load(runner, "toy.symm.upper.2.bg2", "bg2", 2, extra_args)
assert result.exit_code == 0
pixels1 = cooler.Cooler(op.join(datadir, "toy.symm.upper.2.cool")).pixels()[:]
pixels2 = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels2.columns and types.is_float_dtype(
pixels2.dtypes["count"]
)
assert np.allclose(pixels1["count"][:], pixels2["count"][:])
def test_load_field2():
runner = CliRunner()
with runner.isolated_filesystem():
extra_args = ["--count-as-float"]
result = _run_load(runner, "toy.symm.upper.2.bg2", "bg2", 2, extra_args)
assert result.exit_code == 0
pixels1 = cooler.Cooler(op.join(datadir, "toy.symm.upper.2.cool")).pixels()[:]
pixels2 = cooler.Cooler("toy.2.cool").pixels()[:]
assert "count" in pixels2.columns and types.is_float_dtype(
pixels2.dtypes["count"]
)
assert np.allclose(pixels1["count"][:], pixels2["count"][:])
