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

3dmol viewer [NOT TO BE MERGED] #1014

Closed
wants to merge 17 commits into from
Closed
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
4 changes: 4 additions & 0 deletions frontend/3dmol.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
declare module '3dmol/build/3Dmol.js' {
const $3Dmol: any;
export default $3Dmol;
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,17 @@ import { addExperimentThunk, addExperimentWithCheckoutThunk, AppDispatch, experi
import { createExperiment } from "@/lib/redux/slices/experimentAddSlice/asyncActions";

import { DynamicArrayField } from "./DynamicArrayField";
import ThreeDMolViewer from "./ThreeDMolViewer";
import { generateDefaultValues, generateSchema } from "./formGenerator";
import { groupInputs, transformJson } from "./formUtils";

interface FormValues {
name: string;
model: string;
pdbData?: string;
selectedResidues?: string[];
}

export default function NewExperimentForm({ task }: { task: any }) {
const dispatch = useDispatch<AppDispatch>();
const router = useRouter();
Expand All @@ -30,6 +38,9 @@ export default function NewExperimentForm({ task }: { task: any }) {
const walletAddress = user?.wallet?.address;
const userTier = useSelector(selectUserTier);
const isUserSubscribed = useSelector(selectIsUserSubscribed);
const handleViewerSubmit = (data: { pdb: File | null; binderLength: number; hotspots: string }) => {
console.log('Data from viewer:', data);
};

useEffect(() => {
dispatch(refreshUserDataThunk());
Expand All @@ -49,12 +60,12 @@ export default function NewExperimentForm({ task }: { task: any }) {

const groupedInputs = groupInputs(model.ModelJson?.inputs);
const formSchema = generateSchema(model.ModelJson?.inputs);
const defaultValues = generateDefaultValues(model.ModelJson?.inputs, task, model);
// const defaultValues = generateDefaultValues(model.ModelJson?.inputs, task, model);

const form = useForm<z.infer<typeof formSchema>>({
resolver: zodResolver(formSchema),
defaultValues: defaultValues,
});
const form = useForm<FormValues>({
resolver: zodResolver(formSchema), // Make sure formSchema reflects these types
defaultValues: generateDefaultValues(model.ModelJson?.inputs, task, model),
});

// Watch all form values
const watchedValues = form.watch();
Expand Down Expand Up @@ -84,7 +95,9 @@ export default function NewExperimentForm({ task }: { task: any }) {
if (userTier === 'Free') {
// For free tier users, directly create the experiment regardless of subscription status
const response = await dispatch(addExperimentThunk(transformedPayload)).unwrap();
if (response && response.ID) {
if (response.redirectUrl) {
router.push(response.redirectUrl);
} else if (response && response.ID) {
console.log("Experiment created", response);
router.push(`/experiments/${response.ID}`, { scroll: false });
dispatch(experimentListThunk(walletAddress));
Expand Down Expand Up @@ -145,6 +158,12 @@ export default function NewExperimentForm({ task }: { task: any }) {
</div>
</CardContent>
</Card>

<Card>
<CardContent>
<ThreeDMolViewer onSubmit={handleViewerSubmit} />
</CardContent>
</Card>
{!!groupedInputs?.standard && (
<>
<Card>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,9 @@ export default function RerunExperimentForm() {
try {
if (userTier === 'Free' || (userTier === 'Paid' && isUserSubscribed)) {
const response = await addJobToExperiment(experimentID, transformedPayload);
if (response && response.ID) {
if (response.redirectUrl) {
router.push(response.redirectUrl);
} else if (response && response.ID) {
console.log("Job added to experiment", response);
dispatch(experimentDetailThunk(response.ID));
dispatch(experimentListThunk(walletAddress));
Expand Down
195 changes: 195 additions & 0 deletions frontend/app/experiments/(experiment)/(forms)/ThreeDMolViewer.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
import React, { useState, useEffect, useRef } from 'react';
import axios from 'axios';

interface AtomSpec {
chain: string;
resi: number;
}

interface ThreeDMolViewerProps {
onSubmit: (data: { pdb: File | null; binderLength: number; hotspots: string }) => void;
}

const ThreeDMolViewer: React.FC<ThreeDMolViewerProps> = ({ onSubmit }) => {
const [viewer, setViewer] = useState<any>(null);
const [selectedResidues, setSelectedResidues] = useState<Record<string, boolean>>({});
const [binderLength, setBinderLength] = useState(90);
const fileInput = useRef<HTMLInputElement>(null);
const [data, setData] = useState<any>(null);

useEffect(() => {
(async () => {
const module = await import('3dmol/build/3Dmol.js');
const $3Dmol = module.default ? module.default : module;
const element = document.getElementById('container-01') as HTMLElement;
const config = { backgroundColor: 'white' };
const viewer = $3Dmol.createViewer(element, config);
setViewer(viewer);
})();
}, []);

useEffect(() => {
if (viewer) {
updateStyles();
}
}, [selectedResidues]);

const handleFileUpload = (event: React.ChangeEvent<HTMLInputElement>) => {
const file = event.target.files?.[0];
if (file && viewer) {
const reader = new FileReader();
reader.onload = function(e) {
const result = e.target?.result as string;
setData(result);
viewer.addModel(result, "pdb");
updateStyles();
viewer.getModel().setClickable({}, true, (atom: AtomSpec) => {
const residueKey = `${atom.chain}${atom.resi}`;
setSelectedResidues(prev => {
const newResidues = { ...prev };
if (newResidues[residueKey]) {
delete newResidues[residueKey]; // Remove the residue from selected if it's already selected
} else {
newResidues[residueKey] = true; // Add the residue as selected if it's not already
}
return newResidues;
});
});
viewer.render();
};
reader.readAsText(file);
}
};

const updateStyles = () => {
if (!viewer) return;

const allChains = new Set();
const chainsWithSelectedResidues = new Set();

// First, reset styles for all residues
viewer.selectedAtoms({}).forEach((atom: { chain: unknown; resi: any; }) => {
allChains.add(atom.chain);
const residueKey = `${atom.chain}${atom.resi}`;
if (selectedResidues[residueKey]) {
chainsWithSelectedResidues.add(atom.chain);
}
});

if (Object.keys(selectedResidues).length === 0) {
// No residues selected, set all chains to fully visible
viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 1.0 } });
} else {
// Set default style for all chains with higher transparency
viewer.setStyle({}, { cartoon: { color: 'grey', opacity: 0.2 } });

// Set style for chains with selected residues to fully visible
chainsWithSelectedResidues.forEach(chain => {
viewer.setStyle({ chain }, { cartoon: { color: 'grey', opacity: 1.0 } });
});

// Set style for selected residues
Object.keys(selectedResidues).forEach(residueKey => {
const chain = residueKey[0];
const resi = residueKey.slice(1);
viewer.setStyle({ chain, resi }, { cartoon: { color: 'red', opacity: 1.0 } });
});
}

viewer.render();
};

useEffect(() => {
if (viewer) {
const loadAndMakeClickable = async () => {
const data = await axios.get('https://files.rcsb.org/download/1UBQ.pdb').then(res => res.data);
viewer.addModel(data, "pdb");
viewer.getModel().setClickable({}, (atom: { chain: any; resi: any; }) => {
const key = `${atom.chain}:${atom.resi}`;
setSelectedResidues(prev => ({
...prev,
[key]: !prev[key]
}));
});
updateStyles();
};
loadAndMakeClickable();
}
}, [viewer]);



const formatSelectedResidues = () => {
return Object.keys(selectedResidues).map((residueKey) => {
const chain = residueKey[0];
const index = residueKey.slice(1);
return `${chain}${index}`;
}).join(', ');
};

const handleBinderLengthChange = (event: React.ChangeEvent<HTMLInputElement>) => {
setBinderLength(Number(event.target.value));
};

const handleSubmit = async () => {
const pdbFile = fileInput.current?.files?.[0] ?? null;
const hotspots = formatSelectedResidues();
onSubmit({ pdb: pdbFile, binderLength, hotspots });
};

const handleResidueInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
const input = event.target.value;
const newSelectedResidues: Record<string, boolean> = {};

const residues = input.split(',').map(res => res.trim()).filter(res => res);

residues.forEach(residue => {
const match = residue.match(/^([a-zA-Z]+)(\d+)$/);

if (match) {
const chain = match[1].toUpperCase();
const resi = parseInt(match[2], 10);
if (!isNaN(resi)) {
const key = `${chain}${resi}`;
newSelectedResidues[key] = true;
}
}
});

setSelectedResidues(newSelectedResidues);
updateStyles();
};

return (
<div className="relative">
<div style={{
display: 'flex',
justifyContent: 'center',
alignItems: 'center',
padding: '15px',
margin: 'auto',
width: '953px',
height: '377px',
boxSizing: 'border-box' // Includes padding in the width and height
}}>
<div id="container-01" style={{ width: '100%', height: '100%' }}></div>
</div>
<input type="file" ref={fileInput} onChange={handleFileUpload} />
<input type="text" value={formatSelectedResidues()} onChange={handleResidueInputChange} readOnly={false} />
<div>
<label>Binder Length:</label>
<input
type="range"
min="60"
max="120"
value={binderLength}
onChange={handleBinderLengthChange}
/>
<span style={{ marginLeft: '10px' }}>{binderLength}</span>
<button onClick={handleSubmit}>Submit</button>
</div>
</div>
);
};

export default ThreeDMolViewer;
Loading
Loading