Code refactored

Minor code changes.