Merge pull request #19 from AutoResearch/make-return-value-consistent

Make return value consistent
AutoResearch · Jul 26, 2024 · 9a80b1b · 9a80b1b
2 parents 7b533b9 + d00cfdb
commit 9a80b1b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 11 deletions.
diff --git a/docs/Basic Usage.ipynb b/docs/Basic Usage.ipynb
@@ -111,7 +111,7 @@
    },
    "outputs": [],
    "source": [
-    "X_sampled = summed_inequality_sample(condition_pool = X_prime, reference_conditions = X, n = n)\n",
+    "X_sampled = summed_inequality_sample(conditions=X_prime, reference_conditions=X, num_samples=n)\n",
     "print(X_sampled)"
    ]
   },
@@ -160,9 +160,18 @@
    },
    "outputs": [],
    "source": [
-    "X_sampled = summed_inequality_sample(condition_pool = X_prime, reference_conditions = X, n = n)\n",
+    "X_sampled = summed_inequality_sample(conditions=X_prime, reference_conditions=X, num_samples=n)\n",
     "print(X_sampled)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
   }
  ],
  "metadata": {

diff --git a/src/autora/experimentalist/inequality/__init__.py b/src/autora/experimentalist/inequality/__init__.py
@@ -60,28 +60,33 @@ def sample(
     Examples:
         The value 1 is not in the reference. Therefore it is choosen.
         >>> summed_inequality_sample([1, 2, 3], [2, 3, 4])
-        array([[1]])
+           0
+        0  1
 
         The equality distance is set to 0.4. 1 and 1.3 are considered equal, so are 3 and 3.1.
         Therefore 2 is choosen.
         >>> summed_inequality_sample([1, 2, 3], [1.3, 2.7, 3.1], 1, .4)
-        array([[2]])
+           0
+        0  2
 
         The value 3 appears least often in the reference.
         >>> summed_inequality_sample([1, 2, 3], [1, 1, 1, 2, 2, 2, 3, 3])
-        array([[3]])
+           0
+        0  3
 
         The experimentalist "fills up" the reference array so the values are contributed evenly
         >>> summed_inequality_sample([1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 1, 2, 2, 2, 2, 3, 3, 3], 3)
-        array([[1],
-               [3],
-               [1]])
+           0
+        0  1
+        1  3
+        2  1
 
         The experimentalist samples without replacemnt!
         >>> summed_inequality_sample([1, 2, 3], [1, 1, 1], 3)
-        array([[3],
-               [2],
-               [1]])
+           0
+        0  3
+        1  2
+        2  1
 
     """
 
@@ -150,6 +155,8 @@ def sample(
     new_conditions = np.array(condition_pool_res[:num_samples])
     if isinstance(conditions, pd.DataFrame):
         new_conditions = pd.DataFrame(new_conditions, columns=conditions.columns)
+    else:
+        new_conditions = pd.DataFrame(new_conditions)
     return new_conditions