From 41f2e7be57076e8d44b1f6879826b2d31695324c Mon Sep 17 00:00:00 2001
From: Nick Tolley <nic.m.tolley@gmail.com>
Date: Thu, 26 Aug 2021 12:20:26 -0400
Subject: [PATCH] Start system for warning invalid connections

---
 hnn_core/network.py            | 31 +++++++++++++++++++++++++++----
 hnn_core/tests/test_network.py | 22 ++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/hnn_core/network.py b/hnn_core/network.py
index 4269a79cc..01e4f48c3 100644
--- a/hnn_core/network.py
+++ b/hnn_core/network.py
@@ -11,6 +11,7 @@
 from copy import deepcopy
 
 import numpy as np
+from numpy.testing._private.utils import raises
 
 from .drives import _drive_cell_event_times
 from .drives import _get_target_properties, _add_drives_from_params
@@ -1043,12 +1044,34 @@ def add_connection(self, src_gids, target_gids, loc, receptor,
         _validate_type(loc, str, 'loc')
         _validate_type(receptor, str, 'receptor')
 
-        valid_loc = ['proximal', 'distal', 'soma']
-        _check_option('loc', loc, valid_loc)
+        valid_loc = list(self.cell_types[target_type].p_secs.keys())
+        for item in self.cell_types[target_type].sect_loc.values():
+            valid_loc.extend(item)
+
+        valid_loc = list(set(valid_loc))
+        # Separation ecessary to identify valid receptors
+        if loc in self.cell_types[target_type].sect_loc and \
+                self.cell_types[target_type].sect_loc[loc]:
+            sect_list = self.cell_types[target_type].sect_loc[loc]
+        elif loc in self.cell_types[target_type].p_secs:
+            sect_list = [loc]
+        else:
+            raise ValueError(
+                f"The loc '{loc}' is not defined for '{target_type}' "
+                f"cells. Valid locations include {valid_loc}")
+
         conn['loc'] = loc
 
-        valid_receptor = ['ampa', 'nmda', 'gabaa', 'gabab']
-        _check_option('receptor', receptor, valid_receptor)
+        valid_receptors = set(np.concatenate([
+            self.cell_types[target_type].p_secs[sect_name]['syns'] for
+            sect_name in sect_list]))
+
+        if receptor not in valid_receptors:
+            raise ValueError(
+                f"The receptor'{receptor}' is not defined for the loc '{loc}'"
+                f" on '{target_type}' cells. Valid receptors include"
+                f"{valid_receptors}")
+
         conn['receptor'] = receptor
 
         # Create and validate nc_dict
diff --git a/hnn_core/tests/test_network.py b/hnn_core/tests/test_network.py
index bf530a481..c727ac83f 100644
--- a/hnn_core/tests/test_network.py
+++ b/hnn_core/tests/test_network.py
@@ -314,6 +314,28 @@ def test_network():
             kwargs[arg] = string_arg
             net.add_connection(**kwargs)
 
+    # Test warnings for undefined loc
+    with pytest.raises(
+        ValueError,
+        match=r"The loc \'distal\' is not defined for L5_basket"
+              r"cells. Valid locations include \[\'soma\'\]"):
+        kwargs = kwargs_default.copy()
+        kwargs['target_gids'] = 'L5_basket'
+        kwargs['loc'] = 'distal'
+        net.add_connection(**kwargs)
+
+    # Test warnings for undefined synapse for specific loc
+    with pytest.raises(
+        ValueError,
+        match=r"The receptor \'ampa\' is not defined for the loc \'soma\'"
+              r" on \'L5_pyramidal\' cells. Valid receptors include "
+              r"\[\'gabaa\', \'gabab\'\]"):
+        kwargs = kwargs_default.copy()
+        kwargs['target_gids'] = 'L5_pyramidal'
+        kwargs['loc'] = 'soma'
+        kwargs['receptor'] = 'ampa'
+        net.add_connection(**kwargs)
+
     # Check probability=0.5 produces half as many connections as default
     net.add_connection(**kwargs_default)
     kwargs = kwargs_default.copy()