diff --git a/Editor/BrainEditorWindow.cs b/Editor/BrainEditorWindow.cs index 2b16ffe..078095c 100644 --- a/Editor/BrainEditorWindow.cs +++ b/Editor/BrainEditorWindow.cs @@ -12,6 +12,7 @@ namespace NanoBrain { public string title; public Vector2 position; public float radius = 20f; // circle radius + public Nucleus nucleus; } [System.Serializable] @@ -19,15 +20,15 @@ namespace NanoBrain { public int fromId; public int toId; } + public class Dag { + public List nodes = new(); + public List edges = new(); + } public class BrainEditorWindow : EditorWindow { - readonly List nodes = new(); - readonly List edges = new(); + Dag dag = new(); Vector2 pan = Vector2.zero; - float zoom = 1.0f; - const float minZoom = 0.5f; - const float maxZoom = 2.0f; private readonly System.Type acceptedType = typeof(ClusterPrefab); @@ -40,8 +41,9 @@ namespace NanoBrain { void OnEnable() { // Register callback so window updates when selection changes Selection.selectionChanged += OnSelectionChanged; - RefreshSelection(); - ComputeLayout(); + dag = RefreshSelection(); + ComputeLayout(dag); + Repaint(); } private void OnDisable() { @@ -49,33 +51,41 @@ namespace NanoBrain { } private void OnSelectionChanged() { - RefreshSelection(); - ComputeLayout(); + dag = RefreshSelection(); + ComputeLayout(dag); Repaint(); } - private void RefreshSelection() { + private Dag RefreshSelection() { ClusterPrefab prefab = Selection.activeObject as ClusterPrefab; - if (prefab != null && acceptedType.IsAssignableFrom(prefab.GetType())) { - GenerateGraph(prefab); - } + if (prefab != null && acceptedType.IsAssignableFrom(prefab.GetType())) + return GenerateGraph(prefab); + else + return new Dag(); } - private void GenerateGraph(ClusterPrefab prefab) { - nodes.Clear(); - edges.Clear(); + public Dag GenerateGraph(ClusterPrefab prefab) { + Dag dag = new(); int ix = 0; foreach (Nucleus nucleus in prefab.nuclei) { - nodes.Add(new DagNode() { id = ix, title = nucleus.name }); + DagNode node = new() { + id = ix, + title = nucleus.name + }; + dag.nodes.Add(node); if (nucleus is Neuron neuron) { foreach (Nucleus receiver in neuron.receivers) { - int receiverIx = prefab.GetNucleusIndex(receiver); - edges.Add(new DagEdge() { fromId = ix, toId = receiverIx }); + DagEdge edge = new() { + fromId = ix, + toId = prefab.GetNucleusIndex(receiver) + }; + dag.edges.Add(edge); } } ix++; } + return dag; } void OnGUI() { @@ -88,28 +98,28 @@ namespace NanoBrain { Vector2 windowCenter = new(position.width / 2f, position.height / 2f); // compute graph bounds center (in graph space) - Rect bounds = GetGraphBounds(); + Rect bounds = GetGraphBounds(dag); Vector2 graphCenter = bounds.center; // compute autoPan that recenters the graph (does not modify node positions) Vector2 autoPan = -graphCenter; // moves graph center to origin // total translation = windowCenter + autoPan + user pan Matrix4x4 oldMatrix = GUI.matrix; - GUI.matrix = Matrix4x4.TRS(windowCenter + autoPan + pan, Quaternion.identity, Vector3.one * zoom) * + GUI.matrix = Matrix4x4.TRS(windowCenter + autoPan + pan, Quaternion.identity, Vector3.one) * Matrix4x4.TRS(-windowCenter, Quaternion.identity, Vector3.one); // Draw edges first - foreach (DagEdge e in edges) { - DagNode from = GetNodeById(e.fromId); - DagNode to = GetNodeById(e.toId); + foreach (DagEdge e in dag.edges) { + DagNode from = GetNodeById(dag, e.fromId); + DagNode to = GetNodeById(dag, e.toId); if (from == null || to == null) continue; DrawEdgeCircleNodes(from, to); } // Draw nodes (circles) - foreach (DagNode n in nodes) + foreach (DagNode n in dag.nodes) DrawNucleus(n); GUI.matrix = oldMatrix; @@ -118,16 +128,6 @@ namespace NanoBrain { void HandleInput() { Event e = Event.current; - // Zoom with scroll - if (e.type == EventType.ScrollWheel) { - float oldZoom = zoom; - float delta = -e.delta.y * 0.01f; - zoom = Mathf.Clamp(zoom + delta, minZoom, maxZoom); - Vector2 mouse = e.mousePosition; - pan += (mouse - new Vector2(position.width / 2, position.height / 2)) * (1 - zoom / oldZoom); - e.Use(); - } - // Pan with middle or right+ctrl drag if (e.type == EventType.MouseDrag && (e.button == 2 || (e.button == 1 && e.control))) { pan += e.delta; @@ -135,12 +135,12 @@ namespace NanoBrain { } } - DagNode GetNodeById(int id) => nodes.FirstOrDefault(x => x.id == id); + public static DagNode GetNodeById(Dag dag, int id) => dag.nodes.FirstOrDefault(x => x.id == id); - void DrawNucleus(DagNode n) { + public static void DrawNucleus(DagNode n) { Vector3 position = n.position; - Handles.color = Color.white * 0.9f; + Handles.color = Color.black * 0.9f; Handles.DrawSolidDisc(n.position, Vector3.forward, n.radius); Handles.color = Color.white; @@ -153,7 +153,7 @@ namespace NanoBrain { Handles.Label(labelPos, n.title, style); } - void DrawEdgeCircleNodes(DagNode from, DagNode to) { + public static void DrawEdgeCircleNodes(DagNode from, DagNode to) { Vector2 a = from.position; Vector2 b = to.position; if (a == b) return; @@ -162,19 +162,10 @@ namespace NanoBrain { Handles.DrawLine(from.position, to.position); } - // Right-to-left layered layout (sources on the right, sinks on the left) - void ComputeLayout() { - // build adjacency and indegree - Dictionary> adjacency = nodes.ToDictionary(n => n.id, n => new List()); - Dictionary indegree = nodes.ToDictionary(n => n.id, n => 0); - foreach (DagEdge edge in edges) { - if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId)) - continue; - adjacency[edge.fromId].Add(edge.toId); - indegree[edge.toId]++; - } - Dictionary outdegree = nodes.ToDictionary(node => node.id, n => 0); - foreach (DagEdge edge in edges) { + public static void ComputeLayout(Dag dag) { + Dictionary> adjacency = dag.nodes.ToDictionary(n => n.id, n => new List()); + Dictionary outdegree = dag.nodes.ToDictionary(node => node.id, n => 0); + foreach (DagEdge edge in dag.edges) { if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId)) continue; adjacency[edge.fromId].Add(edge.toId); @@ -183,9 +174,10 @@ namespace NanoBrain { // Kahn's algorithm to compute topological layers (horizontal layers) // build parent list (reverse adjacency) and parentIndegree = number of children each parent has - Dictionary> parents = nodes.ToDictionary(n => n.id, _ => new List()); - Dictionary childCount = nodes.ToDictionary(n => n.id, _ => 0); - foreach (DagEdge edge in edges) { + Dictionary> parents = dag.nodes.ToDictionary(n => n.id, _ => new List()); + Dictionary childCount = dag.nodes.ToDictionary(n => n.id, _ => 0); + + foreach (DagEdge edge in dag.edges) { if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId)) continue; adjacency[edge.fromId].Add(edge.toId); parents[edge.toId].Add(edge.fromId); // parent of 'to' is 'from' @@ -210,10 +202,9 @@ namespace NanoBrain { } } - // Any unreachable nodes -> assign next layers int maxLayer = layer.Count > 0 ? layer.Values.Max() : 0; - foreach (DagNode node in nodes) { + foreach (DagNode node in dag.nodes) { if (!layer.ContainsKey(node.id)) { maxLayer++; layer[node.id] = maxLayer; @@ -221,7 +212,12 @@ namespace NanoBrain { } // Group nodes by layer (left to right) - List> layers = layer.GroupBy(kv => kv.Value).OrderBy(g => g.Key).Select(g => g.Select(x => x.Key).ToList()).ToList(); + List> layers = + layer. + GroupBy(kv => kv.Value). + OrderBy(g => g.Key). + Select(g => g.Select(x => x.Key).ToList()). + ToList(); // Same code without using Linq // Build layers dictionary: layerIndex -> List nodeIds @@ -247,25 +243,27 @@ namespace NanoBrain { // } float hSpacing = 100f; - float vSpacing = 100f; + float totalHeight = 400f; // Place nodes: x increases with layer index, y spaced within layer for (int layerIx = 0; layerIx < layers.Count; layerIx++) { List nodeList = layers[layerIx]; - float totalHeight = (nodeList.Count - 1) * vSpacing; + float spacing = totalHeight / nodeList.Count; + float margin = 10 + spacing / 2; for (int i = 0; i < nodeList.Count; i++) { int index = nodeList[i]; - DagNode node = GetNodeById(index); + DagNode node = GetNodeById(dag, index); if (node == null) continue; float x = hSpacing + layerIx * hSpacing; - float y = 400 - totalHeight / 2f + i * vSpacing; + //float y = 400 - totalHeight / 2f + i * vSpacing; + float y = margin + i * spacing; // Debug.Log($"({li}, {i}) -> {x}, {y}"); node.position = new Vector2(x, y); } } - Repaint(); + //Repaint(); } static Rect RectUnion(Rect a, Rect b) { @@ -276,12 +274,13 @@ namespace NanoBrain { return Rect.MinMaxRect(xMin, yMin, xMax, yMax); } - Rect GetGraphBounds() { - if (nodes == null || nodes.Count == 0) return new Rect(Vector2.zero, Vector2.one); + Rect GetGraphBounds(Dag dag) { + if (dag.nodes == null || dag.nodes.Count == 0) + return new Rect(Vector2.zero, Vector2.one); Rect bounds = new( - nodes[0].position - Vector2.one * nodes[0].radius, - 2f * nodes[0].radius * Vector2.one); - foreach (var n in nodes) + dag.nodes[0].position - Vector2.one * dag.nodes[0].radius, + 2f * dag.nodes[0].radius * Vector2.one); + foreach (var n in dag.nodes) bounds = RectUnion(bounds, new Rect(n.position - Vector2.one * n.radius, 2f * n.radius * Vector2.one)); return bounds; diff --git a/Editor/ClusterInspector.cs b/Editor/ClusterInspector.cs index c93e228..06d9db2 100644 --- a/Editor/ClusterInspector.cs +++ b/Editor/ClusterInspector.cs @@ -62,8 +62,18 @@ namespace NanoBrain { public class GraphEditor : GraphView { - public GraphEditor(ClusterPrefab prefab) : base(prefab) { + public enum Mode { + Focus, + Full + } + public Mode mode = Mode.Focus; + public GraphEditor(ClusterPrefab prefab) : base(prefab) { + // create an EnumField for Mode + EnumField enumField = new(mode); + enumField.style.width = 80; + enumField.RegisterValueChangedCallback(OnModeChange); + outputContainer.Insert(0, enumField); Button addButton = new(() => OnAddClusterOutput()) { text = "Add" @@ -71,7 +81,22 @@ namespace NanoBrain { outputContainer.Add(addButton); Add(outputContainer); + } + private void OnModeChange(ChangeEvent evt) { + mode = (Mode)evt.newValue; + + Debug.Log("Mode changed to: " + mode); + } + + Nucleus selectedOutput; + protected override void OnOutputChanged(string outputName) { + if (this.currentNucleus.parent != null) + // Get nucleus in the parent instance + this.selectedOutput = this.currentNucleus.parent.GetNucleus(outputName); + else + // Get nucleus in the prefab + this.selectedOutput = this.prefab.GetNucleus(outputName); } void OnAddClusterOutput() { @@ -83,7 +108,6 @@ namespace NanoBrain { this.currentNucleus = newOutput; } - public void SetGraph(GameObject gameObject, Nucleus nucleus, VisualElement inspectorContainer) { this.gameObject = gameObject; //this.cluster = brain; @@ -123,6 +147,58 @@ namespace NanoBrain { inspectorContainer.Add(inspectorIMGUIContainer); } + protected override void DrawGraph() { + if (mode == Mode.Focus) + DrawFocusGraph(); + else + DrawFullGraph(); + } + + protected void DrawFullGraph() { + Dag dag = GenerateGraph(this.prefab); + BrainEditorWindow.ComputeLayout(dag); + // Draw edges + foreach (DagEdge e in dag.edges) { + DagNode from = dag.nodes.FirstOrDefault(x => x.id == e.fromId); + DagNode to = dag.nodes.FirstOrDefault(x => x.id == e.toId); + if (from == null || to == null) + continue; + + Vector2 fromPosition = from.position; + Vector2 toPosition = to.position; + DrawEdge(fromPosition, toPosition); + } + + // Draw nodes + foreach (DagNode n in dag.nodes) + DrawNucleus(n.nucleus, n.position, 1, n.radius); + } + + public Dag GenerateGraph(ClusterPrefab prefab) { + Dag dag = new(); + + int ix = 0; + foreach (Nucleus nucleus in prefab.nuclei) { + DagNode node = new() { + id = ix, + title = nucleus.name, + nucleus = nucleus + }; + dag.nodes.Add(node); + if (nucleus is Neuron neuron) { + foreach (Nucleus receiver in neuron.receivers) { + DagEdge edge = new() { + fromId = ix, + toId = prefab.GetNucleusIndex(receiver) + }; + dag.edges.Add(edge); + } + } + ix++; + } + return dag; + } + #region Inspector private VisualElement inspectorIMGUIContainer; diff --git a/Editor/ClusterViewer.cs b/Editor/ClusterViewer.cs index 00d6132..241abf7 100644 --- a/Editor/ClusterViewer.cs +++ b/Editor/ClusterViewer.cs @@ -38,9 +38,9 @@ namespace NanoBrain { outputContainer = new() { style = { - flexDirection = FlexDirection.Row, - alignItems = Align.Center, - } + flexDirection = FlexDirection.Row, + alignItems = Align.Center, + } }; List names = this.prefab.outputs.Select(output => output.name).ToList(); @@ -59,7 +59,7 @@ namespace NanoBrain { RegisterCallback(evt => Unsubscribe()); } - void OnOutputChanged(string outputName) { + protected virtual void OnOutputChanged(string outputName) { if (this.currentNucleus.parent != null) // Get nucleus in the parent instance this.currentNucleus = this.currentNucleus.parent.GetNucleus(outputName); @@ -177,7 +177,11 @@ namespace NanoBrain { } - private void DrawGraph() { + protected virtual void DrawGraph() { + DrawFocusGraph(); + } + + protected void DrawFocusGraph() { float size = 20; Vector3 position = new(150, 210, 0); @@ -235,9 +239,6 @@ namespace NanoBrain { // Handles.Label(labelPos, receptorName, style); // } // else { - Handles.color = Color.white; - // The selected nucleus highlight ring - Handles.DrawSolidDisc(position, Vector3.forward, size + 2); float maxValue = 1; if (this.currentNucleus is Neuron neuron) maxValue = neuron.outputMagnitude; @@ -249,9 +250,6 @@ namespace NanoBrain { // } } else { - Handles.color = Color.white; - // The selected nucleus highlight ring - Handles.DrawSolidDisc(position, Vector3.forward, size + 2); float maxValue = 1; if (this.currentNucleus is Neuron neuron) maxValue = neuron.outputMagnitude; @@ -301,8 +299,7 @@ namespace NanoBrain { continue; Vector3 pos = new(50, margin + row * spacing, 0.0f); - Handles.color = Color.white; - Handles.DrawLine(parentPos, pos); + DrawEdge(parentPos, pos); DrawNucleus(receiverNucleus, pos, maxValue, size); row++; @@ -321,7 +318,7 @@ namespace NanoBrain { if (synapse.neuron == null) continue; - if (synapse.neuron.parent is Cluster cluster && + if (synapse.neuron.parent is Cluster cluster && cluster.siblingClusters != null && synapse.neuron.parent != nucleus.parent) { if (drawnArrays.Contains(cluster.siblingClusters)) @@ -372,16 +369,16 @@ namespace NanoBrain { if (synapse.neuron.parent != null && synapse.neuron.parent != this.currentNucleus.parent) { // the synapse nucleus is part of a subcluster //DrawNucleus(synapse.neuron.parent, pos, maxValue, size, color); - DrawNucleus(synapse.neuron, pos, maxValue, size, color); + DrawNucleus(synapse.neuron, pos, size, color); } else { - DrawNucleus(synapse.neuron, pos, maxValue, size, color); + DrawNucleus(synapse.neuron, pos, size, color); } row++; } } - private void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size) { + protected void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size) { Color color; if (Application.isPlaying) { float brightness = 0; @@ -391,10 +388,16 @@ namespace NanoBrain { } else color = Color.black; - DrawNucleus(nucleus, position, maxValue, size, color); + DrawNucleus(nucleus, position, size, color); } - private void DrawNucleus(Nucleus nucleus, Vector3 position, float maxValue, float size, Color color) { + protected void DrawNucleus(Nucleus nucleus, Vector3 position, float size, Color color) { + if (nucleus == this.currentNucleus) { + // The selected nucleus highlight ring + Handles.color = Color.white; + Handles.DrawSolidDisc(position, Vector3.forward, size + 2); + } + if (nucleus is MemoryCell) { Handles.color = Color.white; Handles.DrawWireDisc(position + Vector3.right * 10, Vector3.forward, size); @@ -469,18 +472,18 @@ namespace NanoBrain { // This neuron is part of another cluster parentCluster1.name ??= ""; string baseName = ""; - int colonPos = parentCluster1.name.IndexOf(":"); - if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2) - baseName = parentCluster1.name[..colonPos] + "."; - else - baseName = parentCluster1.name + "."; + int colonPos = parentCluster1.name.IndexOf(":"); + if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2) + baseName = parentCluster1.name[..colonPos] + "."; + else + baseName = parentCluster1.name + "."; // if (colonPos > 0 && colonPos < parentCluster1.name.Length - 2) { // // if it is an array, we should not show the :0 of the first element // //baseName = baseName[..colonPos]; // Handles.Label(labelPos, baseName + nucleus.name, style); // } // else - Handles.Label(labelPos, baseName + nucleus.name, style); + Handles.Label(labelPos, baseName + nucleus.name, style); } else { nucleus.name ??= ""; @@ -562,6 +565,25 @@ namespace NanoBrain { } } + protected void DrawEdge(Vector2 from, Vector2 to) { + Handles.color = Color.white; + Handles.DrawLine(from, to); + } + + // protected void DrawNode(Vector2 position, float size) { + // Handles.color = Color.black * 0.9f; + // Handles.DrawSolidDisc(position, Vector3.forward, size); + + // Handles.color = Color.white; + // GUIStyle style = new(EditorStyles.label) { + // alignment = TextAnchor.UpperCenter, + // normal = { textColor = Color.white }, + // fontStyle = FontStyle.Bold, + // }; + // Vector3 labelPos = position - Vector3.down * (size + 10f); // below disc along up axis + // Handles.Label(labelPos, n.title, style); + // } + void OnSceneGUI(SceneView sceneView) { if (this.gameObject != null) { // if (this.currentNucleus is IReceptor receptor) {