trait x env, for A1=None

limix/__init__.py

@@ -13,7 +13,7 @@
 from ._example import file_example
 from ._testit import test
 
-__version__ = "3.0.1"
+__version__ = "3.0.2b0"
 
 
 __all__ = [

limix/qtl/_result/_mt_result.py

@@ -79,7 +79,7 @@ def _h1_dataframe(self):
                     env_name = "env0_" + str(e)
                     eff = effsizes[l, j]
                     eff_se = effsizes_se[l, j]
-                    v = [i, str(trait), "candidate", str(c), env_name, eff, eff_se]
+                    v = [i, None, "candidate", str(c), env_name, eff, eff_se]
                     h1.append(v)
 
         columns = [
@@ -122,7 +122,7 @@ def _h2_dataframe(self):
                     env_name = "env0_" + str(e)
                     eff = effsizes[l, off + j]
                     eff_se = effsizes_se[l, off + j]
-                    v = [i, str(trait), "candidate", str(c), env_name, eff, eff_se]
+                    v = [i, None, "candidate", str(c), env_name, eff, eff_se]
                     h2.append(v)
 
             off = len(envs0)
@@ -131,7 +131,7 @@ def _h2_dataframe(self):
                     env_name = "env1_" + str(e)
                     eff = effsizes[l, off + j]
                     eff_se = effsizes_se[l, off + j]
-                    v = [i, str(trait), "candidate", str(c), env_name, eff, eff_se]
+                    v = [i, None, "candidate", str(c), env_name, eff, eff_se]
                     h2.append(v)
 
         columns = [

limix/qtl/_result/_st_factory.py

@@ -1,5 +1,6 @@
 from ._st_result import STScanResult
 from ._st_simple import STSimpleModelResult
+from ._tuples import VariantResult, Result
 
 
 class STScanResultFactory:
@@ -26,16 +27,16 @@ def _1d_shape(x):
             return x
 
         def _normalize(h):
-            return {
-                "lml": float(h["lml"]),
-                "covariate_effsizes": _1d_shape(h["covariate_effsizes"]),
-                "candidate_effsizes": _1d_shape(h["candidate_effsizes"]),
-                "covariate_effsizes_se": _1d_shape(h["covariate_effsizes_se"]),
-                "candidate_effsizes_se": _1d_shape(h["candidate_effsizes_se"]),
-                "scale": float(h["scale"]),
-            }
-
-        self._tests.append({"idx": cand_idx, "h2": _normalize(h2)})
+            return VariantResult(
+                lml=float(h["lml"]),
+                covariate_effsizes=_1d_shape(h["covariate_effsizes"]),
+                candidate_effsizes=_1d_shape(h["candidate_effsizes"]),
+                covariate_effsizes_se=_1d_shape(h["covariate_effsizes_se"]),
+                candidate_effsizes_se=_1d_shape(h["candidate_effsizes_se"]),
+                scale=float(h["scale"]),
+            )
+
+        self._tests.append(Result(idx=cand_idx, h2=_normalize(h2)))
 
     def create(self):
         return STScanResult(

limix/qtl/_result/_st_result.py

@@ -57,18 +57,18 @@ def _h2_dataframe(self):
 
         h2 = []
         for i, test in enumerate(self._tests):
-            candidates = list(self._candidates[test["idx"]])
+            candidates = list(self._candidates[test.idx])
 
-            effsizes = test["h2"]["covariate_effsizes"]
-            effsizes_se = test["h2"]["covariate_effsizes_se"]
+            effsizes = test.h2.covariate_effsizes
+            effsizes_se = test.h2.covariate_effsizes_se
             for l, c in enumerate(covariates):
                 eff = effsizes[l]
                 eff_se = effsizes_se[l]
                 v = [i, self._trait, "covariate", str(c), eff, eff_se]
                 h2.append(v)
 
-            effsizes = test["h2"]["candidate_effsizes"]
-            effsizes_se = test["h2"]["candidate_effsizes_se"]
+            effsizes = test.h2.candidate_effsizes
+            effsizes_se = test.h2.candidate_effsizes_se
             for l, c in enumerate(candidates):
                 eff = effsizes[l]
                 eff_se = effsizes_se[l]
@@ -91,10 +91,8 @@ def _stats_dataframe(self):
 
         stats = []
         for i, test in enumerate(self._tests):
-            dof20 = test["h2"]["candidate_effsizes"].size
-            stats.append(
-                [i, self._h0.lml, test["h2"]["lml"], dof20, test["h2"]["scale"]]
-            )
+            dof20 = test.h2.candidate_effsizes.size
+            stats.append([i, self._h0.lml, test.h2.lml, dof20, test.h2.scale])
 
         columns = ["test", "lml0", "lml2", "dof20", "scale2"]
         stats = DataFrame(stats, columns=columns)

limix/qtl/_result/_tuples.py

@@ -0,0 +1,15 @@
+from collections import namedtuple
+
+VariantResult = namedtuple(
+    "VariantResult",
+    [
+        "lml",
+        "covariate_effsizes",
+        "candidate_effsizes",
+        "covariate_effsizes_se",
+        "candidate_effsizes_se",
+        "scale",
+    ],
+)
+
+Result = namedtuple("Result", ["idx", "h1", "h2"], defaults={"h1": None, "h2": None})

limix/qtl/_scan.py

@@ -317,19 +317,27 @@ def _single_trait_scan(idx, lik, Y, M, G, QS, verbose):
 
 
 def _multi_trait_scan(idx, lik, Y, M, G, QS, A, A0, A1, verbose):
-    from xarray import concat
+    from xarray import concat, DataArray
     from numpy import eye, asarray, empty
 
     ntraits = Y.shape[1]
 
     if A1 is None:
         A1 = eye(ntraits)
+        A1 = DataArray(A1, dims=["sample", "env"], coords={"env": Y.trait.values})
 
     if A0 is None:
         A0 = empty((ntraits, 0))
+        A0 = DataArray(A0, dims=["sample", "env"], coords={"env": asarray([], str)})
 
     A0 = _asarray(A0, "env0", ["sample", "env"])
+    if "env" not in A0.coords:
+        A0.coords["env"] = [f"env0_{i}" for i in range(A0.shape[1])]
+
     A1 = _asarray(A1, "env1", ["sample", "env"])
+    if "env" not in A1.coords:
+        A1.coords["env"] = [f"env1_{i}" for i in range(A1.shape[1])]
+
     A01 = concat([A0, A1], dim="env")
 
     if lik[0] == "normal":

limix/qtl/test/test_qtl_scan.py

@@ -1,4 +1,5 @@
 import pytest
+from xarray import DataArray
 import scipy.stats as st
 from numpy import (
     argmin,
@@ -14,7 +15,7 @@
     zeros,
 )
 from numpy.random import RandomState
-from numpy.testing import assert_allclose
+from numpy.testing import assert_allclose, assert_array_equal
 from pandas import DataFrame
 
 from limix.qc import normalise_covariance
@@ -147,6 +148,47 @@ def test_qtl_scan_two_hypotheses_mt():
     str(r)
 
 
+def test_qtl_scan_two_hypotheses_mt_A0A1_none():
+    random = RandomState(0)
+    n = 30
+    ntraits = 2
+    ncovariates = 3
+
+    A = random.randn(ntraits, ntraits)
+    A = A @ A.T
+    M = random.randn(n, ncovariates)
+
+    C0 = random.randn(ntraits, ntraits)
+    C0 = C0 @ C0.T
+
+    C1 = random.randn(ntraits, ntraits)
+    C1 = C1 @ C1.T
+
+    G = random.randn(n, 4)
+
+    A1 = eye(ntraits)
+
+    K = random.randn(n, n + 1)
+    K = normalise_covariance(K @ K.T)
+
+    beta = vec(random.randn(ntraits, ncovariates))
+    alpha = vec(random.randn(A1.shape[1], G.shape[1]))
+
+    m = kron(A, M) @ beta + kron(A1, G) @ alpha
+    Y = unvec(mvn(random, m, kron(C0, K) + kron(C1, eye(n))), (n, -1))
+    Y = DataArray(Y, dims=["sample", "trait"], coords={"trait": ["WA", "Cx"]})
+
+    idx = [[0, 1], 2, [3]]
+    r = scan(G, Y, idx=idx, K=K, M=M, A=A, verbose=False)
+    df = r.effsizes["h2"]
+    df = df[df["test"] == 0]
+    assert_array_equal(df["trait"], ["WA"] * 3 + ["Cx"] * 3 + [None] * 4)
+    assert_array_equal(
+        df["env"], [None] * 6 + ["env1_WA", "env1_WA", "env1_Cx", "env1_Cx"]
+    )
+    str(r)
+
+
 def test_qtl_scan_lmm():
     random = RandomState(0)
     nsamples = 50