Added EGNO block

Author: avisquid

pina/egno_data/1.amc

@@ -0,0 +1,262 @@
+# From official implementation of the Equivariant Graph Neural Operator.
+
+import numpy as np
+import matplotlib.pyplot as plt
+from transforms3d.euler import euler2mat
+from mpl_toolkits.mplot3d import Axes3D
+
+
+class Joint:
+    def __init__(self, name, direction, length, axis, dof, limits):
+        """
+    Definition of basic joint. The joint also contains the information of the
+    bone between it's parent joint and itself. Refer
+    [here](https://research.cs.wisc.edu/graphics/Courses/cs-838-1999/Jeff/ASF-AMC.html)
+    for detailed description for asf files.
+    Parameter
+    ---------
+    name: Name of the joint defined in the asf file. There should always be one
+    root joint. String.
+    direction: Default direction of the joint(bone). The motions are all defined
+    based on this default pose.
+    length: Length of the bone.
+    axis: Axis of rotation for the bone.
+    dof: Degree of freedom. Specifies the number of motion channels and in what
+    order they appear in the AMC file.
+    limits: Limits on each of the channels in the dof specification
+    """
+        self.name = name
+        self.direction = np.reshape(direction, [3, 1])
+        self.length = length
+        axis = np.deg2rad(axis)
+        self.C = euler2mat(*axis)
+        self.Cinv = np.linalg.inv(self.C)
+        self.limits = np.zeros([3, 2])
+        for lm, nm in zip(limits, dof):
+            if nm == 'rx':
+                self.limits[0] = lm
+            elif nm == 'ry':
+                self.limits[1] = lm
+            else:
+                self.limits[2] = lm
+        self.parent = None
+        self.children = []
+        self.coordinate = None
+        self.matrix = None
+
+    def set_motion(self, motion):
+        if self.name == 'root':
+            self.coordinate = np.reshape(np.array(motion['root'][:3]), [3, 1])
+            rotation = np.deg2rad(motion['root'][3:])
+            self.matrix = self.C.dot(euler2mat(*rotation)).dot(self.Cinv)
+        else:
+            idx = 0
+            rotation = np.zeros(3)
+            for axis, lm in enumerate(self.limits):
+                if not np.array_equal(lm, np.zeros(2)):
+                    rotation[axis] = motion[self.name][idx]
+                    idx += 1
+            rotation = np.deg2rad(rotation)
+            self.matrix = self.parent.matrix.dot(self.C).dot(euler2mat(*rotation)).dot(self.Cinv)
+            self.coordinate = self.parent.coordinate + self.length * self.matrix.dot(self.direction)
+        for child in self.children:
+            child.set_motion(motion)
+
+    def get_name_to_idx(self):
+        joints = self.to_dict()
+        name_to_idx = {}
+        for idx, joint in enumerate(joints.values()):
+            assert joint.name not in name_to_idx
+            name_to_idx[joint.name] = idx
+        self.name_to_idx = name_to_idx
+
+    def output_edges(self):
+        joints = self.to_dict()
+        name_to_idx = self.name_to_idx
+        edges = []
+        for idx, joint in enumerate(joints.values()):
+            child = joint
+            if child.parent is not None:
+                parent = child.parent
+                edges.append([name_to_idx[child.name], name_to_idx[parent.name]])
+        return edges
+
+    def output_coord(self):
+        N = len(self.name_to_idx)
+        X = np.zeros((N, 3))
+        joints = self.to_dict()
+        name_to_idx = self.name_to_idx
+        for idx, joint in enumerate(joints.values()):
+            X[name_to_idx[joint.name]] = joint.coordinate.reshape(-1)
+        return X
+
+    def draw(self):
+        joints = self.to_dict()
+        fig = plt.figure()
+        ax = Axes3D(fig)
+
+        ax.set_xlim3d(-50, 10)
+        ax.set_ylim3d(-20, 40)
+        ax.set_zlim3d(-20, 40)
+
+        xs, ys, zs = [], [], []
+        for joint in joints.values():
+            xs.append(joint.coordinate[0, 0])
+            ys.append(joint.coordinate[1, 0])
+            zs.append(joint.coordinate[2, 0])
+        plt.plot(zs, xs, ys, 'b.')
+
+        for joint in joints.values():
+            child = joint
+            if child.parent is not None:
+                parent = child.parent
+                xs = [child.coordinate[0, 0], parent.coordinate[0, 0]]
+                ys = [child.coordinate[1, 0], parent.coordinate[1, 0]]
+                zs = [child.coordinate[2, 0], parent.coordinate[2, 0]]
+                plt.plot(zs, xs, ys, 'r')
+        plt.show()
+
+    def to_dict(self):
+        ret = {self.name: self}
+        for child in self.children:
+            ret.update(child.to_dict())
+        return ret
+
+    def pretty_print(self):
+        print('===================================')
+        print('joint: %s' % self.name)
+        print('direction:')
+        print(self.direction)
+        print('limits:', self.limits)
+        print('parent:', self.parent)
+        print('children:', self.children)
+
+
+def read_line(stream, idx):
+    if idx >= len(stream):
+        return None, idx
+    line = stream[idx].strip().split()
+    idx += 1
+    return line, idx
+
+
+def parse_asf(file_path):
+    '''read joint data only'''
+    with open(file_path) as f:
+        content = f.read().splitlines()
+
+    for idx, line in enumerate(content):
+        # meta infomation is ignored
+        if line == ':bonedata':
+            content = content[idx + 1:]
+            break
+
+    # read joints
+    joints = {'root': Joint('root', np.zeros(3), 0, np.zeros(3), [], [])}
+    idx = 0
+    while True:
+        # the order of each section is hard-coded
+
+        line, idx = read_line(content, idx)
+
+        if line[0] == ':hierarchy':
+            break
+
+        assert line[0] == 'begin'
+
+        line, idx = read_line(content, idx)
+        assert line[0] == 'id'
+
+        line, idx = read_line(content, idx)
+        assert line[0] == 'name'
+        name = line[1]
+
+        line, idx = read_line(content, idx)
+        assert line[0] == 'direction'
+        direction = np.array([float(axis) for axis in line[1:]])
+
+        # skip length
+        line, idx = read_line(content, idx)
+        assert line[0] == 'length'
+        length = float(line[1])
+
+        line, idx = read_line(content, idx)
+        assert line[0] == 'axis'
+        assert line[4] == 'XYZ'
+
+        axis = np.array([float(axis) for axis in line[1:-1]])
+
+        dof = []
+        limits = []
+
+        line, idx = read_line(content, idx)
+        if line[0] == 'dof':
+            dof = line[1:]
+            for i in range(len(dof)):
+                line, idx = read_line(content, idx)
+                if i == 0:
+                    assert line[0] == 'limits'
+                    line = line[1:]
+                assert len(line) == 2
+                mini = float(line[0][1:])
+                maxi = float(line[1][:-1])
+                limits.append((mini, maxi))
+
+            line, idx = read_line(content, idx)
+
+        assert line[0] == 'end'
+        joints[name] = Joint(
+            name,
+            direction,
+            length,
+            axis,
+            dof,
+            limits
+        )
+
+    # read hierarchy
+    assert line[0] == ':hierarchy'
+
+    line, idx = read_line(content, idx)
+
+    assert line[0] == 'begin'
+
+    while True:
+        line, idx = read_line(content, idx)
+        if line[0] == 'end':
+            break
+        assert len(line) >= 2
+        for joint_name in line[1:]:
+            joints[line[0]].children.append(joints[joint_name])
+        for nm in line[1:]:
+            joints[nm].parent = joints[line[0]]
+
+    return joints
+
+
+def parse_amc(file_path):
+    with open(file_path) as f:
+        content = f.read().splitlines()
+
+    for idx, line in enumerate(content):
+        if line == ':DEGREES':
+            content = content[idx + 1:]
+            break
+
+    frames = []
+    idx = 0
+    line, idx = read_line(content, idx)
+    assert line[0].isnumeric(), line
+    EOF = False
+    while not EOF:
+        joint_degree = {}
+        while True:
+            line, idx = read_line(content, idx)
+            if line is None:
+                EOF = True
+                break
+            if line[0].isnumeric():
+                break
+            joint_degree[line[0]] = [float(deg) for deg in line[1:]]
+        frames.append(joint_degree)
+    return frames

pina/egno_data/amc_parser.py

@@ -0,0 +1,112 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "4c55a527",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pina.problem.zoo import SupervisedProblem\n",
+    "import matplotlib.pyplot as plt\n",
+    "from pina.graph import Graph\n",
+    "import amc_parser as amc\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "b71cb46f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Edges -- comuni a tutti i grafi\n",
+    "joints = amc.parse_asf('run.asf')\n",
+    "joints['root'].get_name_to_idx()\n",
+    "edge_idx = torch.tensor(joints['root'].output_edges()).T\n",
+    "\n",
+    "graphs = []\n",
+    "\n",
+    "# Motions -- diversi per ogni grafo\n",
+    "motions = amc.parse_amc('1.amc')\n",
+    "for time, motion in enumerate(motions):\n",
+    "    joints['root'].set_motion(motion)\n",
+    "    pos = torch.tensor(joints['root'].output_coord(), dtype=torch.float32)\n",
+    "    graph = Graph(pos=pos, edge_index=edge_idx, time=torch.tensor(time, dtype=torch.float32), velocity=torch.zeros_like(pos))\n",
+    "    graphs.append(graph)\n",
+    "\n",
+    "for i in range(len(graphs)-1):\n",
+    "    graphs[i].velocity = graphs[i+1].pos - graphs[i].pos\n",
+    "    graphs[-1].velocity = graphs[-2].velocity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "090078da",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Plot each structure with different colors/markers\n",
+    "for i, graph in enumerate(graphs):\n",
+    "    coords = graphs[i].pos.numpy()\n",
+    "    fig = plt.figure()\n",
+    "    ax = fig.add_subplot(111, projection='3d') \n",
+    "    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2])\n",
+    "    ax.set_xlabel(\"X\")\n",
+    "    ax.set_ylabel(\"Y\")\n",
+    "    ax.set_zlabel(\"Z\")\n",
+    "\n",
+    "    ax.set_xlim(-5, 5)\n",
+    "    ax.set_ylim(0, 35)\n",
+    "    ax.set_zlim(0, 45)\n",
+    "\n",
+    "    # Draw edges\n",
+    "    for k, j in graph.edge_index.T.numpy():\n",
+    "        x = [coords[k, 0], coords[j, 0]]\n",
+    "        y = [coords[k, 1], coords[j, 1]]\n",
+    "        z = [coords[k, 2], coords[j, 2]]\n",
+    "        ax.plot(x, y, z, c=\"k\")\n",
+    "        ax.view_init(elev=0, azim=0)\n",
+    "        ax.set_axis_off()\n",
+    "\n",
+    "    plt.savefig(f\"walk_{i}\")\n",
+    "    plt.close()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "ecd595d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input = graphs[:-1]\n",
+    "target = graphs[1:]\n",
+    "problem = SupervisedProblem(input, target)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "deep",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}