Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor tensor containers #222

Merged
merged 4 commits into from
Jan 28, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions algovivo/Muscles.js
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ class Muscles {
if (ten == null) throw new Error("ten required");
this.ten = ten;

this.muscles = null;
this.indices = null;
this.k = Math.fround(90);
this.l0 = null;
this.a = null;
Expand All @@ -19,8 +19,8 @@ class Muscles {
}

get numMuscles() {
if (this.muscles == null) return 0;
return this.muscles.u32().length / 2;
if (this.indices == null) return 0;
return this.indices.u32().length / 2;
}

set(args = {}) {
Expand All @@ -37,8 +37,8 @@ class Muscles {
if (args.k != null) this.k = args.k;

const muscles = mgr.malloc32(numMuscles * 2);
if (this.muscles != null) this.muscles.free();
this.muscles = muscles;
if (this.indices != null) this.indices.free();
this.indices = muscles;

const musclesU32 = muscles.u32();
indices.forEach((m, i) => {
Expand All @@ -59,7 +59,7 @@ class Muscles {
this.numVertices,
args.pos.ptr,
numMuscles,
this.muscles.ptr,
this.indices.ptr,
this.l0.ptr
);
} else {
Expand Down Expand Up @@ -93,10 +93,21 @@ class Muscles {
}
}

toStepArgs() {
const numMuscles = this.numMuscles;
return [
numMuscles,
numMuscles == 0 ? 0 : this.indices.ptr,
this.k,
numMuscles == 0 ? 0 : this.a.ptr,
numMuscles == 0 ? 0 : this.l0.ptr
];
}

dispose() {
if (this.muscles != null) {
this.muscles.free();
this.muscles = null;
if (this.indices != null) {
this.indices.free();
this.indices = null;
}
if (this.l0 != null) {
this.l0.dispose();
Expand Down
99 changes: 42 additions & 57 deletions algovivo/System.js
Original file line number Diff line number Diff line change
Expand Up @@ -25,15 +25,20 @@ class System {

this.spaceDim = args.spaceDim ?? 2;

this._vertices = new Vertices({ ten: this.ten, vertexMass: args.vertexMass, spaceDim: this.spaceDim });
this._muscles = new Muscles({ ten: this.ten });
this._triangles = new Triangles({ ten: this.ten, simplexOrder: this.spaceDim + 1 });
this.vertices = new Vertices({
ten: this.ten,
vertexMass: args.vertexMass,
spaceDim: this.spaceDim
});

this.friction = { k: Math.fround(300) }
}
this.muscles = new Muscles({ ten: this.ten });

this.triangles = new Triangles({
ten: this.ten,
simplexOrder: this.spaceDim + 1
});

get vertices() {
return this._vertices;
this.friction = { k: Math.fround(300) }
}

set fixedVertexId(value) {
Expand All @@ -53,39 +58,31 @@ class System {
}

get vertexMass() {
return this._vertices.vertexMass;
}

get triangles() {
return this._triangles.triangles;
}

set triangles(value) {
this._triangles.triangles = value;
return this.vertices.vertexMass;
}

get rsi() {
return this._triangles.rsi;
return this.triangles.rsi;
}

set rsi(value) {
this._triangles.rsi = value;
this.triangles.rsi = value;
}

get k() {
return this._muscles.k;
return this.muscles.k;
}

set k(value) {
this._muscles.k = value;
this.muscles.k = value;
}

get pos0() {
return this._vertices.pos;
return this.vertices.pos;
}

get vel0() {
return this._vertices.vel0;
return this.vertices.vel0;
}

get pos() {
Expand All @@ -97,58 +94,50 @@ class System {
}

get numVertices() {
return this._vertices.numVertices;
return this.vertices.numVertices;
}

get numTriangles() {
return this._triangles.numTriangles;
return this.triangles.numTriangles;
}

get numMuscles() {
return this._muscles.numMuscles;
}

get muscles() {
return this._muscles.muscles;
}

set muscles(value) {
this._muscles.muscles = value;
return this.muscles.numMuscles;
}

get a() {
return this._muscles.a;
return this.muscles.a;
}

set a(value) {
this._muscles.a = value;
this.muscles.a = value;
}

get l0() {
return this._muscles.l0;
return this.muscles.l0;
}

set l0(value) {
this._muscles.l0 = value;
this.muscles.l0 = value;
}

setVertices(pos) {
this._vertices.set(pos);
this.vertices.set(pos);
}

setMuscles(args = {}) {
this._muscles.set({ ...args, pos: args.pos ?? this.pos0 });
this.muscles.set({ ...args, pos: args.pos ?? this.pos0 });
}

setTriangles(args = {}) {
this._triangles.set({ ...args, pos: args.pos ?? this.pos0 });
this.triangles.set({ ...args, pos: args.pos ?? this.pos0 });
}

getMusclesArray() {
if (this.muscles == null) return [];

const numMuscles = this.numMuscles;
const musclesU32 = this.muscles.u32();
const musclesU32 = this.muscles.indices.u32();
const muscles = [];
for (let i = 0; i < numMuscles; i++) {
const offset = i * 2;
Expand All @@ -164,7 +153,7 @@ class System {
if (this.triangles == null) return [];

const numTriangles = this.numTriangles;
const trianglesU32 = this.triangles.u32();
const trianglesU32 = this.triangles.indices.u32();
const triangles = [];
for (let i = 0; i < numTriangles; i++) {
const offset = i * 3;
Expand Down Expand Up @@ -209,34 +198,30 @@ class System {
numVertices == 0 ? 0 : this.vel0.ptr,
vertexMass,

numMuscles,
numMuscles == 0 ? 0 : this.muscles.ptr,
this.k,
numMuscles == 0 ? 0 : this.a.ptr,
numMuscles == 0 ? 0 : this.l0.ptr,
...this.muscles.toStepArgs(),

...this._triangles.toStepArgs(),
...this.triangles.toStepArgs(),

this.friction.k,

fixedVertexId,

numVertices == 0 ? 0 : this._vertices.pos1.ptr,
numVertices == 0 ? 0 : this._vertices.posGrad.ptr,
numVertices == 0 ? 0 : this._vertices.posTmp.ptr,
numVertices == 0 ? 0 : this._vertices.vel1.ptr,
numVertices == 0 ? 0 : this.vertices.pos1.ptr,
numVertices == 0 ? 0 : this.vertices.posGrad.ptr,
numVertices == 0 ? 0 : this.vertices.posTmp.ptr,
numVertices == 0 ? 0 : this.vertices.vel1.ptr,
);

if (numVertices != 0) {
this._vertices.pos0.slot.f32().set(this._vertices.pos1.slot.f32());
this._vertices.vel0.slot.f32().set(this._vertices.vel1.slot.f32());
this.vertices.pos0.slot.f32().set(this.vertices.pos1.slot.f32());
this.vertices.vel0.slot.f32().set(this.vertices.vel1.slot.f32());
}
}

dispose() {
this._vertices.dispose();
this._muscles.dispose();
this._triangles.dispose();
this.vertices.dispose();
this.muscles.dispose();
this.triangles.dispose();
}
}

Expand Down
20 changes: 8 additions & 12 deletions algovivo/Triangles.js
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ class Triangles {
this.ten = ten;

this.simplexOrder = args.simplexOrder ?? 3;
this.triangles = null;
this.indices = null;
this.rsi = null;
this.mu = null;
this.lambda = null;
Expand All @@ -28,10 +28,6 @@ class Triangles {
return this.numElements;
}

get indices() {
return this.triangles;
}

toStepArgs() {
const numElements = this.numElements;
return [
Expand All @@ -55,9 +51,9 @@ class Triangles {
const mgr = this.memoryManager;
const ten = this.ten;

const triangles = indices ? mgr.malloc32(numTriangles * this.simplexOrder) : this.triangles;
if (indices && this.triangles != null) this.triangles.free();
this.triangles = triangles;
const triangles = indices ? mgr.malloc32(numTriangles * this.simplexOrder) : this.indices;
if (indices && this.indices != null) this.indices.free();
this.indices = triangles;

if (indices != null) {
const trianglesU32 = triangles.u32();
Expand Down Expand Up @@ -89,7 +85,7 @@ class Triangles {
this.numVertices,
pos.ptr,
numTriangles,
this.triangles.ptr,
this.indices.ptr,
this.rsi.ptr
);

Expand All @@ -108,9 +104,9 @@ class Triangles {
}

dispose() {
if (this.triangles != null) {
this.triangles.free();
this.triangles = null;
if (this.indices != null) {
this.indices.free();
this.indices = null;
}
if (this.rsi != null) {
this.rsi.dispose();
Expand Down
Loading